Star887
1791
代码
代码
Issues252
Pull Requests305
Actions
项目讨论
Wiki
分析
Star887
1791
已合并
refactor: rename APIs in ops-nn batch4 (10 ops) #8569
hahaha22创建于 19 天前
refactor: rename APIs in ops-nn batch4 (10 ops) #8569
已合并
hahaha22创建于 19 天前
从hahaha22/ops-nnrename-api-batch4合入到cann/ops-nnmaster
共 38 个文件变更+4908-4905
Mactivation/softmax_cross_entropy_with_logits/op_kernel/arch35/softmax_cross_entropy_with_logits_full_load.h+38-38
@@ -25,9 +25,9 @@ using namespace AscendC;
25using AscendC::MicroAPI::LoadDist;25using AscendC::MicroAPI::LoadDist;
26using AscendC::MicroAPI::MaskPattern;26using AscendC::MicroAPI::MaskPattern;
27using AscendC::MicroAPI::MaskReg;27using AscendC::MicroAPI::MaskReg;
28-using AscendC::MicroAPI::MaskUnPack;
29using AscendC::MicroAPI::RegTensor;28using AscendC::MicroAPI::RegTensor;
30using AscendC::MicroAPI::StoreDist;29using AscendC::MicroAPI::StoreDist;
30+using AscendC::MicroAPI::UnPack;
31using AscendC::MicroAPI::UpdateMask;31using AscendC::MicroAPI::UpdateMask;
32 32 
33template <typename T, uint64_t schId, uint64_t featuresBrc, uint64_t labelsBrc, uint64_t db>33template <typename T, uint64_t schId, uint64_t featuresBrc, uint64_t labelsBrc, uint64_t db>
@@ -283,8 +283,8 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsFullLoad<T, schId, featuresB
283 if constexpr (IsSameType<T, half>::value) {283 if constexpr (IsSameType<T, half>::value) {
284 constValue = -65504;284 constValue = -65504;
285 }285 }
286- static constexpr MultiCopyConfig config = {false};286+ static constexpr NdDmaConfig config = {false};
287- MultiCopyLoopInfo<2> loopInfo;287+ NdDmaLoopInfo<2> loopInfo;
288 loopInfo.loopSize[0] = r_;288 loopInfo.loopSize[0] = r_;
289 loopInfo.loopSize[1] = tileNum;289 loopInfo.loopSize[1] = tileNum;
290 loopInfo.loopLpSize[0] = 0;290 loopInfo.loopLpSize[0] = 0;
@@ -302,7 +302,7 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsFullLoad<T, schId, featuresB
302 }302 }
303 loopInfo.loopDstStride[0] = 1;303 loopInfo.loopDstStride[0] = 1;
304 loopInfo.loopDstStride[1] = rNumAlign;304 loopInfo.loopDstStride[1] = rNumAlign;
305- MultiCopyParams<T, 2> paramsMain = {loopInfo, constValue};305+ NdDmaParams<T, 2> paramsMain = {loopInfo, constValue};
306 DataCopy<T, 2, config>(dstBuf, srcTensor[offset], paramsMain);306 DataCopy<T, 2, config>(dstBuf, srcTensor[offset], paramsMain);
307 dstQueue.EnQue<T>(dstBuf);307 dstQueue.EnQue<T>(dstBuf);
308}308}
@@ -373,12 +373,12 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsFullLoad<T, schId, featuresB
373 373 
374 for (uint16_t i = 0; i < aTimes; i++) {374 for (uint16_t i = 0; i < aTimes; i++) {
375 AscendC::MicroAPI::Duplicate(featuresReg, minValue);375 AscendC::MicroAPI::Duplicate(featuresReg, minValue);
376- AscendC::MicroAPI::DataCopy(featuresReg1, featuresAddr + i * rAlign + repeatTimes * vfLen);376+ AscendC::MicroAPI::LoadAlign(featuresReg1, featuresAddr + i * rAlign + repeatTimes * vfLen);
377 AscendC::MicroAPI::Max(featuresReg1, featuresReg, featuresReg1, preg);377 AscendC::MicroAPI::Max(featuresReg1, featuresReg, featuresReg1, preg);
378- AscendC::MicroAPI::Copy<T, AscendC::MicroAPI::MaskMergeMode::MERGING>(featuresReg, featuresReg1, preg);378+ AscendC::MicroAPI::Move<T, AscendC::MicroAPI::MaskMergeMode::MERGING>(featuresReg, featuresReg1, preg);
379 for (uint16_t j = 0; j < repeatTimes; j++) {379 for (uint16_t j = 0; j < repeatTimes; j++) {
380 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<T>(i, rAlign, j, vfLen);380 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<T>(i, rAlign, j, vfLen);
381- AscendC::MicroAPI::DataCopy(featuresReg1, featuresAddr1, offset);381+ AscendC::MicroAPI::LoadAlign(featuresReg1, featuresAddr1, offset);
382 AscendC::MicroAPI::Max(featuresReg, featuresReg1, featuresReg, pregMain);382 AscendC::MicroAPI::Max(featuresReg, featuresReg1, featuresReg, pregMain);
383 }383 }
384 if constexpr (sizeof(T) == 2) {384 if constexpr (sizeof(T) == 2) {
@@ -391,11 +391,11 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsFullLoad<T, schId, featuresB
391 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresRegLowest32, featuresRegLowest, pregReduce);391 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresRegLowest32, featuresRegLowest, pregReduce);
392 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresRegHighest32, featuresRegHighest, pregReduce);392 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresRegHighest32, featuresRegHighest, pregReduce);
393 AscendC::MicroAPI::Max(maxRegTemp, featuresRegLowest32, featuresRegHighest32, pregReduce);393 AscendC::MicroAPI::Max(maxRegTemp, featuresRegLowest32, featuresRegHighest32, pregReduce);
394- AscendC::MicroAPI::ReduceMax(maxReg, maxRegTemp, pregReduce);394+ AscendC::MicroAPI::Reduce<ReduceType::MAX>(maxReg, maxRegTemp, pregReduce);
395 } else {395 } else {
396- AscendC::MicroAPI::ReduceMax(maxReg, featuresReg, pregReduce);396+ AscendC::MicroAPI::Reduce<ReduceType::MAX>(maxReg, featuresReg, pregReduce);
397 }397 }
398- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(maxAddr + i, maxReg, mergePreg);398+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(maxAddr + i, maxReg, mergePreg);
399 }399 }
400 }400 }
401}401}
@@ -432,35 +432,35 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsFullLoad<T, schId, featuresB
432 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);432 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
433 433 
434 for (uint16_t i = 0; i < aTimes; i++) {434 for (uint16_t i = 0; i < aTimes; i++) {
435- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg32, maxAddr + i);435+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg32, maxAddr + i);
436 for (uint16_t j = 0; j < repeatTimes; j++) {436 for (uint16_t j = 0; j < repeatTimes; j++) {
437 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T>(i, rAlign, j, vfLen);437 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T>(i, rAlign, j, vfLen);
438 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);438 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);
439 if constexpr (sizeof(T) == 2) {439 if constexpr (sizeof(T) == 2) {
440- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(featuresReg,440+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
441- featuresAddr, offsetT);441+ featuresReg, featuresAddr, offsetT);
442 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresReg32, featuresReg, pregMain);442 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresReg32, featuresReg, pregMain);
443 } else {443 } else {
444- AscendC::MicroAPI::DataCopy(featuresReg32, featuresAddr, offset);444+ AscendC::MicroAPI::LoadAlign(featuresReg32, featuresAddr, offset);
445 }445 }
446 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, pregMain);446 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, pregMain);
447 AscendC::MicroAPI::Exp(temp1Reg, subReg, pregMain);447 AscendC::MicroAPI::Exp(temp1Reg, subReg, pregMain);
448- AscendC::MicroAPI::DataCopy(temp1Addr, temp1Reg, offset, pregMain);448+ AscendC::MicroAPI::StoreAlign(temp1Addr, temp1Reg, offset, pregMain);
449- AscendC::MicroAPI::DataCopy(subAddr, subReg, offset, pregMain);449+ AscendC::MicroAPI::StoreAlign(subAddr, subReg, offset, pregMain);
450 }450 }
451 451 
452 for (uint16_t k = 0; k < tailLoop; k++) {452 for (uint16_t k = 0; k < tailLoop; k++) {
453 if constexpr (sizeof(T) == 2) {453 if constexpr (sizeof(T) == 2) {
454- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(454+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
455 featuresReg, featuresAddr + i * rAlign + repeatTimes * vfLen);455 featuresReg, featuresAddr + i * rAlign + repeatTimes * vfLen);
456 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresReg32, featuresReg, preg);456 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(featuresReg32, featuresReg, preg);
457 } else {457 } else {
458- AscendC::MicroAPI::DataCopy(featuresReg32, featuresAddr + i * rAlign + repeatTimes * vfLen);458+ AscendC::MicroAPI::LoadAlign(featuresReg32, featuresAddr + i * rAlign + repeatTimes * vfLen);
459 }459 }
460 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, preg);460 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, preg);
461 AscendC::MicroAPI::Exp(temp1Reg, subReg, preg);461 AscendC::MicroAPI::Exp(temp1Reg, subReg, preg);
462- AscendC::MicroAPI::DataCopy(temp1Addr + i * rAlign + repeatTimes * vfLen, temp1Reg, pregAlign);462+ AscendC::MicroAPI::StoreAlign(temp1Addr + i * rAlign + repeatTimes * vfLen, temp1Reg, pregAlign);
463- AscendC::MicroAPI::DataCopy(subAddr + i * rAlign + repeatTimes * vfLen, subReg, preg);463+ AscendC::MicroAPI::StoreAlign(subAddr + i * rAlign + repeatTimes * vfLen, subReg, preg);
464 }464 }
465 }465 }
466 }466 }
@@ -506,63 +506,63 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsFullLoad<T, schId, featuresB
506 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);506 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
507 507 
508 for (uint16_t i = 0; i < aTimes; i++) {508 for (uint16_t i = 0; i < aTimes; i++) {
509- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(sumReg, sumAddr + i);509+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(sumReg, sumAddr + i);
510 for (uint16_t j = 0; j < repeatTimes; j++) {510 for (uint16_t j = 0; j < repeatTimes; j++) {
511 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T>(i, rAlign, j, vfLen);511 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T>(i, rAlign, j, vfLen);
512 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);512 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);
513- AscendC::MicroAPI::DataCopy(temp1Reg, temp1Addr, offset);513+ AscendC::MicroAPI::LoadAlign(temp1Reg, temp1Addr, offset);
514 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, pregMain);514 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, pregMain);
515 if constexpr (sizeof(T) == 2) {515 if constexpr (sizeof(T) == 2) {
516- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(labelsReg, labelsAddr,516+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(labelsReg, labelsAddr,
517- offsetT);517+ offsetT);
518 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(labelsReg32, labelsReg, pregMain);518 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(labelsReg32, labelsReg, pregMain);
519 } else {519 } else {
520- AscendC::MicroAPI::DataCopy(labelsReg32, labelsAddr, offset);520+ AscendC::MicroAPI::LoadAlign(labelsReg32, labelsAddr, offset);
521 }521 }
522 AscendC::MicroAPI::Sub(backPropReg32, temp1Reg, labelsReg32, pregMain);522 AscendC::MicroAPI::Sub(backPropReg32, temp1Reg, labelsReg32, pregMain);
523 if constexpr (sizeof(T) == 2) {523 if constexpr (sizeof(T) == 2) {
524 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, backPropReg32, pregMain);524 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, backPropReg32, pregMain);
525- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(525+ AscendC::MicroAPI::StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(
526 backPropAddr, backPropReg, offsetT, pregMain);526 backPropAddr, backPropReg, offsetT, pregMain);
527 } else {527 } else {
528- AscendC::MicroAPI::DataCopy(backPropAddr, backPropReg32, offset, pregMain);528+ AscendC::MicroAPI::StoreAlign(backPropAddr, backPropReg32, offset, pregMain);
529 }529 }
530- AscendC::MicroAPI::DataCopy(subReg, subAddr, offset);530+ AscendC::MicroAPI::LoadAlign(subReg, subAddr, offset);
531 AscendC::MicroAPI::Log(logReg, sumReg, pregMain);531 AscendC::MicroAPI::Log(logReg, sumReg, pregMain);
532 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, pregMain);532 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, pregMain);
533 AscendC::MicroAPI::Mul(temp2Reg, temp2Reg, labelsReg32, pregMain);533 AscendC::MicroAPI::Mul(temp2Reg, temp2Reg, labelsReg32, pregMain);
534- AscendC::MicroAPI::DataCopy(temp2Addr, temp2Reg, offset, pregMain);534+ AscendC::MicroAPI::StoreAlign(temp2Addr, temp2Reg, offset, pregMain);
535 }535 }
536 536 
537 for (uint16_t k = 0; k < tailLoop; k++) {537 for (uint16_t k = 0; k < tailLoop; k++) {
538- AscendC::MicroAPI::DataCopy(temp1Reg, temp1Addr + i * rAlign + repeatTimes * vfLen);538+ AscendC::MicroAPI::LoadAlign(temp1Reg, temp1Addr + i * rAlign + repeatTimes * vfLen);
539 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, preg);539 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, preg);
540 if constexpr (sizeof(T) == 2) {540 if constexpr (sizeof(T) == 2) {
541- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(541+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
542 labelsReg, labelsAddr + i * rAlign + repeatTimes * vfLen);542 labelsReg, labelsAddr + i * rAlign + repeatTimes * vfLen);
543 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(labelsReg32, labelsReg, preg);543 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(labelsReg32, labelsReg, preg);
544 } else {544 } else {
545- AscendC::MicroAPI::DataCopy(labelsReg32, labelsAddr + i * rAlign + repeatTimes * vfLen);545+ AscendC::MicroAPI::LoadAlign(labelsReg32, labelsAddr + i * rAlign + repeatTimes * vfLen);
546 }546 }
547 AscendC::MicroAPI::Sub(backPropReg32, temp1Reg, labelsReg32, preg);547 AscendC::MicroAPI::Sub(backPropReg32, temp1Reg, labelsReg32, preg);
548 548 
549 if constexpr (sizeof(T) == 2) {549 if constexpr (sizeof(T) == 2) {
550 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, backPropReg32, preg);550 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, backPropReg32, preg);
551- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(551+ AscendC::MicroAPI::StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(
552 backPropAddr + i * rAlign + repeatTimes * vfLen, backPropReg, preg);552 backPropAddr + i * rAlign + repeatTimes * vfLen, backPropReg, preg);
553 } else {553 } else {
554- AscendC::MicroAPI::DataCopy(backPropAddr + i * rAlign + repeatTimes * vfLen, backPropReg32,554+ AscendC::MicroAPI::StoreAlign(backPropAddr + i * rAlign + repeatTimes * vfLen, backPropReg32,
555- pregAlign);555+ pregAlign);
556 }556 }
557- AscendC::MicroAPI::DataCopy(subReg, subAddr + i * rAlign + repeatTimes * vfLen);557+ AscendC::MicroAPI::LoadAlign(subReg, subAddr + i * rAlign + repeatTimes * vfLen);
558 AscendC::MicroAPI::Log(logReg, sumReg, preg);558 AscendC::MicroAPI::Log(logReg, sumReg, preg);
559 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, preg);559 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, preg);
560 AscendC::MicroAPI::Mul(temp2Reg, temp2Reg, labelsReg32, preg);560 AscendC::MicroAPI::Mul(temp2Reg, temp2Reg, labelsReg32, preg);
561- AscendC::MicroAPI::DataCopy(temp2Addr + i * rAlign + repeatTimes * vfLen, temp2Reg, pregAlign);561+ AscendC::MicroAPI::StoreAlign(temp2Addr + i * rAlign + repeatTimes * vfLen, temp2Reg, pregAlign);
562 }562 }
563 }563 }
564 }564 }
565}565}
566 566 
567} // namespace SoftmaxCrossEntropyWithLogits567} // namespace SoftmaxCrossEntropyWithLogits
568-#endif568+#endif
Mactivation/softmax_cross_entropy_with_logits/op_kernel/arch35/softmax_cross_entropy_with_logits_split_r.h+112-110
@@ -107,14 +107,14 @@ private:
107 int64_t tailBlockFactor; // a轴分核，尾核数据量107 int64_t tailBlockFactor; // a轴分核，尾核数据量
108 int64_t rUbNumFactor; // R轴切分，一次UB可以放下的数据量，全载模板下等于r，注意32b对齐108 int64_t rUbNumFactor; // R轴切分，一次UB可以放下的数据量，全载模板下等于r，注意32b对齐
109 int64_t aUbNumFactor; // A轴切分，一次UB可以放下的数据量，非全载模板下等于1，注意32b对齐109 int64_t aUbNumFactor; // A轴切分，一次UB可以放下的数据量，非全载模板下等于1，注意32b对齐
110- int64_t aLoopTimes; //主核A方向循环搬移数据的次数110+ int64_t aLoopTimes; // 主核A方向循环搬移数据的次数
111- int64_t aLoopTimesT; //尾核A方向循环搬移数据的次数111+ int64_t aLoopTimesT; // 尾核A方向循环搬移数据的次数
112- int64_t aLoopTail; //主核A方向尾块的数据量112+ int64_t aLoopTail; // 主核A方向尾块的数据量
113- int64_t aLoopTailT; //尾核A方向尾块的数据量113+ int64_t aLoopTailT; // 尾核A方向尾块的数据量
114- int64_t rLoopTime; //不能全载时，R轴反向的循环次数114+ int64_t rLoopTime; // 不能全载时，R轴反向的循环次数
115- int64_t rLoopTile; //不能全载时，R轴反向的尾块数据量115+ int64_t rLoopTile; // 不能全载时，R轴反向的尾块数据量
116- int64_t kTimesTail; //不能全载时，完全二分累加，存在主尾块相加的次数116+ int64_t kTimesTail; // 不能全载时，完全二分累加，存在主尾块相加的次数
117- int64_t kTimes; //不能全载时，完全二分累加，2的k次方内循环次数117+ int64_t kTimes; // 不能全载时，完全二分累加，2的k次方内循环次数
118 int64_t updateStart_;118 int64_t updateStart_;
119 int64_t rLoopTileAlign;119 int64_t rLoopTileAlign;
120 int64_t featuresBrcDim;120 int64_t featuresBrcDim;
@@ -207,8 +207,8 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
207{207{
208 LocalTensor<T> xInUb = featuresQueue_.AllocTensor<T>();208 LocalTensor<T> xInUb = featuresQueue_.AllocTensor<T>();
209 T constValue = padValue;209 T constValue = padValue;
210- static constexpr MultiCopyConfig config = {false};210+ static constexpr NdDmaConfig config = {false};
211- MultiCopyLoopInfo<2> loopInfo;211+ NdDmaLoopInfo<2> loopInfo;
212 212 
213 loopInfo.loopSize[0] = onceR;213 loopInfo.loopSize[0] = onceR;
214 loopInfo.loopSize[1] = nTailNum;214 loopInfo.loopSize[1] = nTailNum;
@@ -226,7 +226,7 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
226 }226 }
227 loopInfo.loopDstStride[0] = 1;227 loopInfo.loopDstStride[0] = 1;
228 loopInfo.loopDstStride[1] = onceRAlign;228 loopInfo.loopDstStride[1] = onceRAlign;
229- MultiCopyParams<T, 2> paramsMain = {loopInfo, constValue};229+ NdDmaParams<T, 2> paramsMain = {loopInfo, constValue};
230 DataCopy<T, 2, config>(xInUb, featuresGm_[offset], paramsMain);230 DataCopy<T, 2, config>(xInUb, featuresGm_[offset], paramsMain);
231 featuresQueue_.EnQue(xInUb);231 featuresQueue_.EnQue(xInUb);
232}232}
@@ -237,8 +237,8 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
237{237{
238 LocalTensor<T> xInUb = labelsQueue_.AllocTensor<T>();238 LocalTensor<T> xInUb = labelsQueue_.AllocTensor<T>();
239 T constValue = padValue;239 T constValue = padValue;
240- static constexpr MultiCopyConfig config = {false};240+ static constexpr NdDmaConfig config = {false};
241- MultiCopyLoopInfo<2> loopInfo;241+ NdDmaLoopInfo<2> loopInfo;
242 242 
243 loopInfo.loopSize[0] = onceR;243 loopInfo.loopSize[0] = onceR;
244 loopInfo.loopSize[1] = nTailNum;244 loopInfo.loopSize[1] = nTailNum;
@@ -256,7 +256,7 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
256 }256 }
257 loopInfo.loopDstStride[0] = 1;257 loopInfo.loopDstStride[0] = 1;
258 loopInfo.loopDstStride[1] = onceRAlign;258 loopInfo.loopDstStride[1] = onceRAlign;
259- MultiCopyParams<T, 2> paramsMain = {loopInfo, constValue};259+ NdDmaParams<T, 2> paramsMain = {loopInfo, constValue};
260 DataCopy<T, 2, config>(xInUb, labelsGm_[offset], paramsMain);260 DataCopy<T, 2, config>(xInUb, labelsGm_[offset], paramsMain);
261 labelsQueue_.EnQue(xInUb);261 labelsQueue_.EnQue(xInUb);
262}262}
@@ -349,17 +349,17 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
349 for (uint16_t j = 0; j < repeatTimes1; j++) {349 for (uint16_t j = 0; j < repeatTimes1; j++) {
350 preg = AscendC::MicroAPI::UpdateMask<float>(nTimes);350 preg = AscendC::MicroAPI::UpdateMask<float>(nTimes);
351 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLenfp32);351 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLenfp32);
352- AscendC::MicroAPI::DataCopy(srcReg1, maxUbAddr, srcOffset);352+ AscendC::MicroAPI::LoadAlign(srcReg1, maxUbAddr, srcOffset);
353 if constexpr (sizeof(T) == 4) {353 if constexpr (sizeof(T) == 4) {
354- AscendC::MicroAPI::DataCopy(srcReg2, maxUbOnceAddrB32, srcOffset);354+ AscendC::MicroAPI::LoadAlign(srcReg2, maxUbOnceAddrB32, srcOffset);
355 } else {355 } else {
356 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<T>(j, vfLenfp32);356 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<T>(j, vfLenfp32);
357- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg2B16, maxUbOnceAddr,357+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg2B16, maxUbOnceAddr,
358- srcOffset1);358+ srcOffset1);
359 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg2, srcReg2B16, preg);359 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg2, srcReg2B16, preg);
360 }360 }
361 AscendC::MicroAPI::Max(srcReg3, srcReg1, srcReg2, preg);361 AscendC::MicroAPI::Max(srcReg3, srcReg1, srcReg2, preg);
362- AscendC::MicroAPI::DataCopy(maxUbAddr, srcReg3, srcOffset, preg);362+ AscendC::MicroAPI::StoreAlign(maxUbAddr, srcReg3, srcOffset, preg);
363 }363 }
364 }364 }
365}365}
@@ -411,7 +411,7 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
411 preg = AscendC::MicroAPI::UpdateMask<float>(nSize);411 preg = AscendC::MicroAPI::UpdateMask<float>(nSize);
412 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);412 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);
413 AscendC::MicroAPI::Duplicate(srcReg0, minValueFp32_);413 AscendC::MicroAPI::Duplicate(srcReg0, minValueFp32_);
414- AscendC::MicroAPI::DataCopy(maxUbAddr, srcReg0, srcOffset, preg);414+ AscendC::MicroAPI::StoreAlign(maxUbAddr, srcReg0, srcOffset, preg);
415 }415 }
416 }416 }
417}417}
@@ -440,14 +440,14 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
440 for (uint16_t i = 0; i < outerLoopTimes; ++i) {440 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
441 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);441 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
442 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat);442 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat);
443- AscendC::MicroAPI::DataCopy(aReg, srcUbAddr, srcOffset);443+ AscendC::MicroAPI::LoadAlign(aReg, srcUbAddr, srcOffset);
444 for (uint16_t j = 0; j < innerLoopTimes; ++j) {444 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
445 AscendC::MicroAPI::AddrReg srcOffsetJ = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat, j,445 AscendC::MicroAPI::AddrReg srcOffsetJ = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat, j,
446 stride);446 stride);
447- AscendC::MicroAPI::DataCopy(bReg, dstUbAddr, srcOffsetJ);447+ AscendC::MicroAPI::LoadAlign(bReg, dstUbAddr, srcOffsetJ);
448 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);448 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
449 }449 }
450- AscendC::MicroAPI::DataCopy(cahUbAddr, aReg, srcOffset, pMask);450+ AscendC::MicroAPI::StoreAlign(cahUbAddr, aReg, srcOffset, pMask);
451 }451 }
452 }452 }
453}453}
@@ -477,24 +477,24 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
477 AscendC::MicroAPI::RegTensor<float> subReg;477 AscendC::MicroAPI::RegTensor<float> subReg;
478 AscendC::MicroAPI::RegTensor<float> expReg;478 AscendC::MicroAPI::RegTensor<float> expReg;
479 for (uint16_t i = 0; i < nTimes; i++) {479 for (uint16_t i = 0; i < nTimes; i++) {
480- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(480+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
481 maxReg,481 maxReg,
482 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast482 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast
483 for (uint16_t j = 0; j < repeatTimes1; j++) {483 for (uint16_t j = 0; j < repeatTimes1; j++) {
484 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j, vfLen);484 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j, vfLen);
485 if constexpr (sizeof(T) == 2) {485 if constexpr (sizeof(T) == 2) {
486 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNum, j, vfLen);486 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNum, j, vfLen);
487- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,487+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
488- srcOffset);488+ srcOffset);
489 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);489 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
490 } else {490 } else {
491 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j,491 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j,
492 vfLen);492 vfLen);
493- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset1);493+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset1);
494 }494 }
495 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);495 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
496 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);496 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
497- AscendC::MicroAPI::DataCopy(outUbAddr, expReg, outOffset, copyOutReg);497+ AscendC::MicroAPI::StoreAlign(outUbAddr, expReg, outOffset, copyOutReg);
498 }498 }
499 }499 }
500 }500 }
@@ -532,37 +532,37 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
532 AscendC::MicroAPI::RegTensor<float> outReg1;532 AscendC::MicroAPI::RegTensor<float> outReg1;
533 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);533 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
534 for (uint16_t i = 0; i < nTimes; i++) {534 for (uint16_t i = 0; i < nTimes; i++) {
535- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);535+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);
536 for (uint16_t j = 0; j < repeatTimes1; j++) {536 for (uint16_t j = 0; j < repeatTimes1; j++) {
537 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNumTAlign, j, vfLen);537 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNumTAlign, j, vfLen);
538 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,538 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
539 vfLen);539 vfLen);
540 if constexpr (sizeof(T) == 2) {540 if constexpr (sizeof(T) == 2) {
541- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,541+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
542- srcOffset);542+ srcOffset);
543 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);543 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
544 } else {544 } else {
545- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);545+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
546 }546 }
547 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);547 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
548 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);548 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
549- AscendC::MicroAPI::DataCopy(outReg, outUbAddr, outOffset);549+ AscendC::MicroAPI::LoadAlign(outReg, outUbAddr, outOffset);
550 AscendC::MicroAPI::Add(outReg1, expReg, outReg, copyOutReg);550 AscendC::MicroAPI::Add(outReg1, expReg, outReg, copyOutReg);
551- AscendC::MicroAPI::DataCopy(outUbAddr, outReg1, outOffset, copyOutReg);551+ AscendC::MicroAPI::StoreAlign(outUbAddr, outReg1, outOffset, copyOutReg);
552 }552 }
553 for (uint16_t k = 0; k < tailLoopTimes; k++) {553 for (uint16_t k = 0; k < tailLoopTimes; k++) {
554 if constexpr (sizeof(T) == 2) {554 if constexpr (sizeof(T) == 2) {
555- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(555+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
556 srcReg0, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);556 srcReg0, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
557 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);557 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);
558 } else {558 } else {
559- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);559+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
560 }560 }
561 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);561 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
562 AscendC::MicroAPI::Exp(expReg, subReg, preg);562 AscendC::MicroAPI::Exp(expReg, subReg, preg);
563- AscendC::MicroAPI::DataCopy(outReg, outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);563+ AscendC::MicroAPI::LoadAlign(outReg, outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
564 AscendC::MicroAPI::Add(outReg1, expReg, outReg, preg);564 AscendC::MicroAPI::Add(outReg1, expReg, outReg, preg);
565- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, outReg1, preg);565+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, outReg1, preg);
566 }566 }
567 }567 }
568 }568 }
@@ -600,7 +600,7 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
600 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);600 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
601 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);601 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
602 for (uint16_t i = 0; i < nTimes; i++) {602 for (uint16_t i = 0; i < nTimes; i++) {
603- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(603+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
604 maxReg,604 maxReg,
605 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast605 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast
606 for (uint16_t j = 0; j < repeatTimes1; j++) {606 for (uint16_t j = 0; j < repeatTimes1; j++) {
@@ -608,27 +608,27 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
608 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,608 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
609 vfLen);609 vfLen);
610 if constexpr (sizeof(T) == 2) {610 if constexpr (sizeof(T) == 2) {
611- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,611+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
612- srcOffset);612+ srcOffset);
613 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);613 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
614 } else {614 } else {
615- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);615+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
616 }616 }
617 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);617 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
618 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);618 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
619- AscendC::MicroAPI::DataCopy(outUbAddr, expReg, outOffset, copyOutReg);619+ AscendC::MicroAPI::StoreAlign(outUbAddr, expReg, outOffset, copyOutReg);
620 }620 }
621 for (uint16_t k = 0; k < tailLoopTimes; k++) {621 for (uint16_t k = 0; k < tailLoopTimes; k++) {
622 if constexpr (sizeof(T) == 2) {622 if constexpr (sizeof(T) == 2) {
623- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(623+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
624 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);624 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
625 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);625 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);
626 } else {626 } else {
627- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);627+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
628 }628 }
629 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);629 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
630 AscendC::MicroAPI::Exp(expReg, subReg, preg);630 AscendC::MicroAPI::Exp(expReg, subReg, preg);
631- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg, preg1);631+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg, preg1);
632 }632 }
633 }633 }
634 }634 }
@@ -721,11 +721,11 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
721 for (uint16_t j = 0; j < repeatTimes1; j++) {721 for (uint16_t j = 0; j < repeatTimes1; j++) {
722 preg = AscendC::MicroAPI::UpdateMask<float>(nSize1);722 preg = AscendC::MicroAPI::UpdateMask<float>(nSize1);
723 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);723 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);
724- AscendC::MicroAPI::DataCopy(srcReg0, cacheUbAddr, srcOffset);724+ AscendC::MicroAPI::LoadAlign(srcReg0, cacheUbAddr, srcOffset);
725 AscendC::MicroAPI::Log(logReg0, srcReg0, preg);725 AscendC::MicroAPI::Log(logReg0, srcReg0, preg);
726- AscendC::MicroAPI::Copy(sumReg0, srcReg0, preg);726+ AscendC::MicroAPI::Move(sumReg0, srcReg0, preg);
727- AscendC::MicroAPI::DataCopy(logUbAddr, logReg0, srcOffset, preg);727+ AscendC::MicroAPI::StoreAlign(logUbAddr, logReg0, srcOffset, preg);
728- AscendC::MicroAPI::DataCopy(sumUbAddr, sumReg0, srcOffset, preg);728+ AscendC::MicroAPI::StoreAlign(sumUbAddr, sumReg0, srcOffset, preg);
729 }729 }
730 }730 }
731}731}
@@ -771,28 +771,28 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
771 771 
772 AscendC::MicroAPI::RegTensor<T> backProbReg;772 AscendC::MicroAPI::RegTensor<T> backProbReg;
773 for (uint16_t i = 0; i < nTimes; i++) {773 for (uint16_t i = 0; i < nTimes; i++) {
774- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(774+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
775 maxReg,775 maxReg,
776 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast776 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast
777- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(777+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
778 sumReg,778 sumReg,
779 sumUbAddr + i); // mask通过data_copy的brc功能实现broadcast779 sumUbAddr + i); // mask通过data_copy的brc功能实现broadcast
780- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(780+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
781 logReg,781 logReg,
782 logUbAddr + i); // mask通过data_copy的brc功能实现broadcast782 logUbAddr + i); // mask通过data_copy的brc功能实现broadcast
783 for (uint16_t j = 0; j < repeatTimes1; j++) {783 for (uint16_t j = 0; j < repeatTimes1; j++) {
784 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNum, j, vfLen);784 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNum, j, vfLen);
785 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j, vfLen);785 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j, vfLen);
786 if constexpr (sizeof(T) == 2) {786 if constexpr (sizeof(T) == 2) {
787- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,787+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
788- srcOffset);788+ srcOffset);
789 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);789 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
790- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg1, labelsUbAddr,790+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg1, labelsUbAddr,
791- srcOffset);791+ srcOffset);
792 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, copyOutReg);792 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, copyOutReg);
793 } else {793 } else {
794- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);794+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
795- AscendC::MicroAPI::DataCopy(srcReg1fp32, labelsUbAddr, srcOffset);795+ AscendC::MicroAPI::LoadAlign(srcReg1fp32, labelsUbAddr, srcOffset);
796 }796 }
797 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);797 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
798 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);798 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
@@ -800,14 +800,14 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
800 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, copyOutReg);800 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, copyOutReg);
801 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);801 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);
802 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, copyOutReg);802 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, copyOutReg);
803- AscendC::MicroAPI::DataCopy(outUbAddr, tmpReg, outOffset, copyOutReg);803+ AscendC::MicroAPI::StoreAlign(outUbAddr, tmpReg, outOffset, copyOutReg);
804 if constexpr (sizeof(T) == 2) {804 if constexpr (sizeof(T) == 2) {
805 AscendC::MicroAPI::AddrReg outOffset1 = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNum, j, vfLen);805 AscendC::MicroAPI::AddrReg outOffset1 = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNum, j, vfLen);
806 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backProbReg, expReg, copyOutReg);806 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backProbReg, expReg, copyOutReg);
807- AscendC::MicroAPI::DataCopy<T, StoreDist::DIST_PACK_B32>(backProbAddr, backProbReg, outOffset1,807+ AscendC::MicroAPI::StoreAlign<T, StoreDist::DIST_PACK_B32>(backProbAddr, backProbReg, outOffset1,
808- copyOutReg);808+ copyOutReg);
809 } else {809 } else {
810- AscendC::MicroAPI::DataCopy(backProbAddr, expReg, outOffset, copyOutReg);810+ AscendC::MicroAPI::StoreAlign(backProbAddr, expReg, outOffset, copyOutReg);
811 }811 }
812 }812 }
813 }813 }
@@ -864,13 +864,13 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
864 AscendC::MicroAPI::RegTensor<T> backPropReg;864 AscendC::MicroAPI::RegTensor<T> backPropReg;
865 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);865 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
866 for (uint16_t i = 0; i < nTimes; i++) {866 for (uint16_t i = 0; i < nTimes; i++) {
867- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(867+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
868 maxReg,868 maxReg,
869 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast869 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast
870- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(870+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
871 sumReg,871 sumReg,
872 sumUbAddr + i); // mask通过data_copy的brc功能实现broadcast872 sumUbAddr + i); // mask通过data_copy的brc功能实现broadcast
873- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(873+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
874 logReg,874 logReg,
875 logUbAddr + i); // mask通过data_copy的brc功能实现broadcast875 logUbAddr + i); // mask通过data_copy的brc功能实现broadcast
876 for (uint16_t j = 0; j < repeatTimes1; j++) {876 for (uint16_t j = 0; j < repeatTimes1; j++) {
@@ -878,15 +878,15 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
878 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,878 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
879 vfLen);879 vfLen);
880 if constexpr (sizeof(T) == 2) {880 if constexpr (sizeof(T) == 2) {
881- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,881+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
882- srcOffset);882+ srcOffset);
883 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);883 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
884- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg1, labelsUbAddr,884+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg1, labelsUbAddr,
885- srcOffset);885+ srcOffset);
886 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, copyOutReg);886 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, copyOutReg);
887 } else {887 } else {
888- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);888+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
889- AscendC::MicroAPI::DataCopy(srcReg1fp32, labelsUbAddr, srcOffset);889+ AscendC::MicroAPI::LoadAlign(srcReg1fp32, labelsUbAddr, srcOffset);
890 }890 }
891 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);891 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
892 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);892 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
@@ -894,28 +894,28 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
894 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, copyOutReg);894 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, copyOutReg);
895 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);895 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);
896 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, copyOutReg);896 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, copyOutReg);
897- AscendC::MicroAPI::DataCopy(outReg, outUbAddr, outOffset);897+ AscendC::MicroAPI::LoadAlign(outReg, outUbAddr, outOffset);
898 AscendC::MicroAPI::Add(outReg1, tmpReg, outReg, copyOutReg);898 AscendC::MicroAPI::Add(outReg1, tmpReg, outReg, copyOutReg);
899- AscendC::MicroAPI::DataCopy(outUbAddr, outReg1, outOffset, copyOutReg);899+ AscendC::MicroAPI::StoreAlign(outUbAddr, outReg1, outOffset, copyOutReg);
900 if constexpr (sizeof(T) == 2) {900 if constexpr (sizeof(T) == 2) {
901 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, expReg, copyOutReg);901 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, expReg, copyOutReg);
902- AscendC::MicroAPI::DataCopy<T, StoreDist::DIST_PACK_B32>(backProbAddr, backPropReg, srcOffset,902+ AscendC::MicroAPI::StoreAlign<T, StoreDist::DIST_PACK_B32>(backProbAddr, backPropReg, srcOffset,
903- copyOutReg);903+ copyOutReg);
904 } else {904 } else {
905- AscendC::MicroAPI::DataCopy(backProbAddr, expReg, srcOffset, copyOutReg);905+ AscendC::MicroAPI::StoreAlign(backProbAddr, expReg, srcOffset, copyOutReg);
906 }906 }
907 }907 }
908 for (uint16_t k = 0; k < tailLoopTimes; k++) {908 for (uint16_t k = 0; k < tailLoopTimes; k++) {
909 if constexpr (sizeof(T) == 2) {909 if constexpr (sizeof(T) == 2) {
910- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(910+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
911 srcReg0, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);911 srcReg0, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
912 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);912 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);
913- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(913+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
914 srcReg1, labelsUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);914 srcReg1, labelsUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
915 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, preg);915 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, preg);
916 } else {916 } else {
917- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);917+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
918- AscendC::MicroAPI::DataCopy(srcReg1fp32, labelsUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);918+ AscendC::MicroAPI::LoadAlign(srcReg1fp32, labelsUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
919 }919 }
920 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);920 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
921 AscendC::MicroAPI::Exp(expReg, subReg, preg);921 AscendC::MicroAPI::Exp(expReg, subReg, preg);
@@ -923,15 +923,16 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
923 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, preg);923 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, preg);
924 AscendC::MicroAPI::Sub(logReg, logReg, subReg, preg);924 AscendC::MicroAPI::Sub(logReg, logReg, subReg, preg);
925 AscendC::MicroAPI::Mul(logReg, logReg, srcReg1fp32, preg);925 AscendC::MicroAPI::Mul(logReg, logReg, srcReg1fp32, preg);
926- AscendC::MicroAPI::DataCopy(outReg, outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);926+ AscendC::MicroAPI::LoadAlign(outReg, outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
927 AscendC::MicroAPI::Add(outReg1, logReg, outReg, preg);927 AscendC::MicroAPI::Add(outReg1, logReg, outReg, preg);
928- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, outReg1, preg);928+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, outReg1, preg);
929 if constexpr (sizeof(T) == 2) {929 if constexpr (sizeof(T) == 2) {
930 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, expReg, preg);930 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backPropReg, expReg, preg);
931- AscendC::MicroAPI::DataCopy<T, StoreDist::DIST_PACK_B32>(931+ AscendC::MicroAPI::StoreAlign<T, StoreDist::DIST_PACK_B32>(
932 backProbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen, backPropReg, preg);932 backProbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen, backPropReg, preg);
933 } else {933 } else {
934- AscendC::MicroAPI::DataCopy(backProbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen, expReg, preg);934+ AscendC::MicroAPI::StoreAlign(backProbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen, expReg,
935+ preg);
935 }936 }
936 }937 }
937 }938 }
@@ -985,13 +986,13 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
985 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);986 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
986 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);987 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
987 for (uint16_t i = 0; i < nTimes; i++) {988 for (uint16_t i = 0; i < nTimes; i++) {
988- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(989+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
989 maxReg,990 maxReg,
990 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast991 maxUbAddr + i); // mask通过data_copy的brc功能实现broadcast
991- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(992+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
992 sumReg,993 sumReg,
993 sumUbAddr + i); // mask通过data_copy的brc功能实现broadcast994 sumUbAddr + i); // mask通过data_copy的brc功能实现broadcast
994- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(995+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(
995 logReg,996 logReg,
996 logUbAddr + i); // mask通过data_copy的brc功能实现broadcast997 logUbAddr + i); // mask通过data_copy的brc功能实现broadcast
997 for (uint16_t j = 0; j < repeatTimes1; j++) {998 for (uint16_t j = 0; j < repeatTimes1; j++) {
@@ -999,15 +1000,15 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
999 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,1000 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
1000 vfLen);1001 vfLen);
1001 if constexpr (sizeof(T) == 2) {1002 if constexpr (sizeof(T) == 2) {
1002- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,1003+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
1003- srcOffset);1004+ srcOffset);
1004 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);1005 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
1005- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg1, labelsUbAddr,1006+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg1, labelsUbAddr,
1006- srcOffset);1007+ srcOffset);
1007 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, copyOutReg);1008 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, copyOutReg);
1008 } else {1009 } else {
1009- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);1010+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
1010- AscendC::MicroAPI::DataCopy(srcReg1fp32, labelsUbAddr, srcOffset);1011+ AscendC::MicroAPI::LoadAlign(srcReg1fp32, labelsUbAddr, srcOffset);
1011 }1012 }
1012 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);1013 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
1013 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);1014 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
@@ -1015,28 +1016,28 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
1015 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, copyOutReg);1016 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, copyOutReg);
1016 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);1017 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);
1017 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, copyOutReg);1018 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, copyOutReg);
1018- AscendC::MicroAPI::DataCopy(outUbAddr, tmpReg, outOffset, copyOutReg);1019+ AscendC::MicroAPI::StoreAlign(outUbAddr, tmpReg, outOffset, copyOutReg);
1019 if constexpr (sizeof(T) == 2) {1020 if constexpr (sizeof(T) == 2) {
1020 AscendC::MicroAPI::AddrReg outOffset1 = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNumAlign, j,1021 AscendC::MicroAPI::AddrReg outOffset1 = AscendC::MicroAPI::CreateAddrReg<T>(i, rOnceNumAlign, j,
1021 vfLen);1022 vfLen);
1022 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backProbReg, expReg, copyOutReg);1023 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backProbReg, expReg, copyOutReg);
1023- AscendC::MicroAPI::DataCopy<T, StoreDist::DIST_PACK_B32>(backProbAddr, backProbReg, outOffset1,1024+ AscendC::MicroAPI::StoreAlign<T, StoreDist::DIST_PACK_B32>(backProbAddr, backProbReg, outOffset1,
1024- copyOutReg);1025+ copyOutReg);
1025 } else {1026 } else {
1026- AscendC::MicroAPI::DataCopy(backProbAddr, expReg, outOffset, copyOutReg);1027+ AscendC::MicroAPI::StoreAlign(backProbAddr, expReg, outOffset, copyOutReg);
1027 }1028 }
1028 }1029 }
1029 for (uint16_t k = 0; k < tailLoopTimes; k++) {1030 for (uint16_t k = 0; k < tailLoopTimes; k++) {
1030 if constexpr (sizeof(T) == 2) {1031 if constexpr (sizeof(T) == 2) {
1031- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(1032+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
1032 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);1033 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
1033 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);1034 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcRegfp32, srcReg0, preg);
1034- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(1035+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
1035 srcReg1, labelsUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);1036 srcReg1, labelsUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
1036 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, preg);1037 AscendC::MicroAPI::Cast<float, T, castB16ToB32>(srcReg1fp32, srcReg1, preg);
1037 } else {1038 } else {
1038- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);1039+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
1039- AscendC::MicroAPI::DataCopy(srcReg1fp32, labelsUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);1040+ AscendC::MicroAPI::LoadAlign(srcReg1fp32, labelsUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
1040 }1041 }
1041 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);1042 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
1042 AscendC::MicroAPI::Exp(expReg, subReg, preg);1043 AscendC::MicroAPI::Exp(expReg, subReg, preg);
@@ -1044,13 +1045,14 @@ SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc, labelsBrc, db>::Compu
1044 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, preg);1045 AscendC::MicroAPI::Sub(expReg, expReg, srcReg1fp32, preg);
1045 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, preg);1046 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, preg);
1046 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, preg);1047 AscendC::MicroAPI::Mul(tmpReg, tmpReg, srcReg1fp32, preg);
1047- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, tmpReg, preg1);1048+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, tmpReg, preg1);
1048 if constexpr (sizeof(T) == 2) {1049 if constexpr (sizeof(T) == 2) {
1049 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backProbReg, expReg, preg);1050 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(backProbReg, expReg, preg);
1050- AscendC::MicroAPI::DataCopy<T, StoreDist::DIST_PACK_B32>(1051+ AscendC::MicroAPI::StoreAlign<T, StoreDist::DIST_PACK_B32>(
1051 backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, backProbReg, preg);1052 backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, backProbReg, preg);
1052 } else {1053 } else {
1053- AscendC::MicroAPI::DataCopy(backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg, preg1);1054+ AscendC::MicroAPI::StoreAlign(backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg,
1055+ preg1);
1054 }1056 }
1055 }1057 }
1056 }1058 }
@@ -1097,13 +1099,13 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
1097 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);1099 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);
1098 AscendC::MicroAPI::MaskReg1100 AscendC::MicroAPI::MaskReg
1099 regAllFp32 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();1101 regAllFp32 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
1100- AscendC::MicroAPI::DataCopy(srcReg0, cacheUbAddr, srcOffset);1102+ AscendC::MicroAPI::LoadAlign(srcReg0, cacheUbAddr, srcOffset);
1101 if constexpr (sizeof(T) == 2) {1103 if constexpr (sizeof(T) == 2) {
1102 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(lossReg0, srcReg0, preg);1104 AscendC::MicroAPI::Cast<T, float, castB32ToB16>(lossReg0, srcReg0, preg);
1103- AscendC::MicroAPI::DataCopy<T, StoreDist::DIST_PACK_B32>(lossUbAddr, lossReg0, preg);1105+ AscendC::MicroAPI::StoreAlign<T, StoreDist::DIST_PACK_B32>(lossUbAddr, lossReg0, preg);
1104 } else {1106 } else {
1105- AscendC::MicroAPI::Copy(lossReg0, srcReg0, preg);1107+ AscendC::MicroAPI::Move(lossReg0, srcReg0, preg);
1106- AscendC::MicroAPI::DataCopy(lossUbAddr, lossReg0, srcOffset, preg);1108+ AscendC::MicroAPI::StoreAlign(lossUbAddr, lossReg0, srcOffset, preg);
1107 }1109 }
1108 }1110 }
1109 }1111 }
@@ -1219,4 +1221,4 @@ __aicore__ inline void SoftmaxCrossEntropyWithLogitsSplitR<T, schId, featuresBrc
1219}1221}
1220 1222 
1221} // namespace SoftmaxCrossEntropyWithLogits1223} // namespace SoftmaxCrossEntropyWithLogits
1222-#endif // SOFTMAX_CROSS_ENTROPY_WITH_LOGITS_SPLIT_R_H1224+#endif // SOFTMAX_CROSS_ENTROPY_WITH_LOGITS_SPLIT_R_H
Mactivation/softmax_grad/op_kernel/arch35/softmax_grad_ar_full_load.h+59-56
@@ -30,6 +30,9 @@ using AscendC::MicroAPI::MaskMergeMode;
30using AscendC::MicroAPI::MaskReg;30using AscendC::MicroAPI::MaskReg;
31using AscendC::MicroAPI::RegTensor;31using AscendC::MicroAPI::RegTensor;
32using AscendC::MicroAPI::StoreDist;32using AscendC::MicroAPI::StoreDist;
33+using AscendC::Reg::LoadAlign;
34+using AscendC::Reg::Reduce;
35+using AscendC::Reg::StoreAlign;
33 36 
34static constexpr uint32_t DOUBLE_BUFFER = 2;37static constexpr uint32_t DOUBLE_BUFFER = 2;
35static constexpr uint32_t BLOCK_SIZE = platform::GetUbBlockSize();38static constexpr uint32_t BLOCK_SIZE = platform::GetUbBlockSize();
@@ -58,10 +61,10 @@ private:
58 61 
59 __aicore__ inline void CopyOutY(int64_t ubA, int64_t offset);62 __aicore__ inline void CopyOutY(int64_t ubA, int64_t offset);
60 63 
61- __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T* dst, RegTensor<float>& src, MaskReg& preg,64+ __aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T* dst, RegTensor<float>& src, MaskReg& preg,
62 uint32_t offset);65 uint32_t offset);
63 66 
64- __aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg,67+ __aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg,
65 uint32_t offset);68 uint32_t offset);
66 69 
67private:70private:
@@ -151,9 +154,9 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputeSmallR(const int64_t aSize)
151 154 
152 uint16_t loopTimes = aSize;155 uint16_t loopTimes = aSize;
153 if (rSize <= VL_FP32) {156 if (rSize <= VL_FP32) {
154- __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr();157+ __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr();
155- __local_mem__ T* x0 = (__local_mem__ T*)x0Tensor.GetPhyAddr();158+ __ubuf__ T* x0 = (__ubuf__ T*)x0Tensor.GetPhyAddr();
156- __local_mem__ T* x1 = (__local_mem__ T*)x1Tensor.GetPhyAddr();159+ __ubuf__ T* x1 = (__ubuf__ T*)x1Tensor.GetPhyAddr();
157 160 
158 __VEC_SCOPE__161 __VEC_SCOPE__
159 {162 {
@@ -168,7 +171,7 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputeSmallR(const int64_t aSize)
168 LoadTensorForDtypeTIn(x1, reg1, pMask, i * rAligned);171 LoadTensorForDtypeTIn(x1, reg1, pMask, i * rAligned);
169 Mul(reg2, reg0, reg1, pMask);172 Mul(reg2, reg0, reg1, pMask);
170 173 
171- ReduceSum(reg2, reg2, pMask);174+ Reduce<ReduceType::SUM>(reg2, reg2, pMask);
172 Duplicate(reg2, reg2, pFull);175 Duplicate(reg2, reg2, pFull);
173 176 
174 Mul(reg1, reg0, reg1, pMask);177 Mul(reg1, reg0, reg1, pMask);
@@ -179,11 +182,11 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputeSmallR(const int64_t aSize)
179 }182 }
180 }183 }
181 } else {184 } else {
182- __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr();185+ __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr();
183- __local_mem__ T* x0 = (__local_mem__ T*)x0Tensor.GetPhyAddr();186+ __ubuf__ T* x0 = (__ubuf__ T*)x0Tensor.GetPhyAddr();
184- __local_mem__ T* x1 = (__local_mem__ T*)x1Tensor.GetPhyAddr();187+ __ubuf__ T* x1 = (__ubuf__ T*)x1Tensor.GetPhyAddr();
185- __local_mem__ T* x0_1 = (__local_mem__ T*)x0Tensor.GetPhyAddr() + VL_FP32;188+ __ubuf__ T* x0_1 = (__ubuf__ T*)x0Tensor.GetPhyAddr() + VL_FP32;
186- __local_mem__ T* x1_1 = (__local_mem__ T*)x1Tensor.GetPhyAddr() + VL_FP32;189+ __ubuf__ T* x1_1 = (__ubuf__ T*)x1Tensor.GetPhyAddr() + VL_FP32;
187 190 
188 __VEC_SCOPE__191 __VEC_SCOPE__
189 {192 {
@@ -203,8 +206,8 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputeSmallR(const int64_t aSize)
203 Mul(reg2_1, reg0_1, reg1_1, pMask);206 Mul(reg2_1, reg0_1, reg1_1, pMask);
204 207 
205 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2_1, reg2, reg2_1, pMask);208 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2_1, reg2, reg2_1, pMask);
206- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg2, reg2_1, pMask);209+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg2, reg2_1, pMask);
207- ReduceSum(reg2, reg2, pFull);210+ Reduce<ReduceType::SUM>(reg2, reg2, pFull);
208 Duplicate(reg2, reg2, pFull);211 Duplicate(reg2, reg2, pFull);
209 212 
210 Mul(reg1, reg0, reg1, pFull);213 Mul(reg1, reg0, reg1, pFull);
@@ -254,26 +257,26 @@ __aicore__ inline void SoftmaxGradAR<T>::NormCompute(const int64_t aSize)
254 int64_t tailSrcBOffset = floorVLCount * VL_FP32;257 int64_t tailSrcBOffset = floorVLCount * VL_FP32;
255 int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32;258 int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32;
256 259 
257- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr();260+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr();
258- __local_mem__ T* foldX0A = (__local_mem__ T*)x0Tensor.GetPhyAddr();261+ __ubuf__ T* foldX0A = (__ubuf__ T*)x0Tensor.GetPhyAddr();
259- __local_mem__ T* foldX0B = (__local_mem__ T*)x0Tensor.GetPhyAddr() + foldSrcBOffset;262+ __ubuf__ T* foldX0B = (__ubuf__ T*)x0Tensor.GetPhyAddr() + foldSrcBOffset;
260- __local_mem__ T* tailX0A = (__local_mem__ T*)x0Tensor.GetPhyAddr() + tailSrcAOffset;263+ __ubuf__ T* tailX0A = (__ubuf__ T*)x0Tensor.GetPhyAddr() + tailSrcAOffset;
261- __local_mem__ T* tailX0B = (__local_mem__ T*)x0Tensor.GetPhyAddr() + tailSrcBOffset;264+ __ubuf__ T* tailX0B = (__ubuf__ T*)x0Tensor.GetPhyAddr() + tailSrcBOffset;
262- __local_mem__ T* unFoldX0 = (__local_mem__ T*)x0Tensor.GetPhyAddr() + unFoldSrcOffset;265+ __ubuf__ T* unFoldX0 = (__ubuf__ T*)x0Tensor.GetPhyAddr() + unFoldSrcOffset;
263 266 
264- __local_mem__ T* foldX1A = (__local_mem__ T*)x1Tensor.GetPhyAddr();267+ __ubuf__ T* foldX1A = (__ubuf__ T*)x1Tensor.GetPhyAddr();
265- __local_mem__ T* foldX1B = (__local_mem__ T*)x1Tensor.GetPhyAddr() + foldSrcBOffset;268+ __ubuf__ T* foldX1B = (__ubuf__ T*)x1Tensor.GetPhyAddr() + foldSrcBOffset;
266- __local_mem__ T* tailX1A = (__local_mem__ T*)x1Tensor.GetPhyAddr() + tailSrcAOffset;269+ __ubuf__ T* tailX1A = (__ubuf__ T*)x1Tensor.GetPhyAddr() + tailSrcAOffset;
267- __local_mem__ T* tailX1B = (__local_mem__ T*)x1Tensor.GetPhyAddr() + tailSrcBOffset;270+ __ubuf__ T* tailX1B = (__ubuf__ T*)x1Tensor.GetPhyAddr() + tailSrcBOffset;
268- __local_mem__ T* unFoldX1 = (__local_mem__ T*)x1Tensor.GetPhyAddr() + unFoldSrcOffset;271+ __ubuf__ T* unFoldX1 = (__ubuf__ T*)x1Tensor.GetPhyAddr() + unFoldSrcOffset;
269 272 
270 __VEC_SCOPE__273 __VEC_SCOPE__
271 {274 {
272 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();275 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
273- AscendC::MicroAPI::UnalignReg UReg;276+ AscendC::MicroAPI::UnalignRegForStore UReg;
274 277 
275 for (uint16_t i = 0; i < outerLoopTimes; i++) {278 for (uint16_t i = 0; i < outerLoopTimes; i++) {
276- dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;279+ dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;
277 for (uint16_t j = 0; j < mainFoldLoopTimes; j++) {280 for (uint16_t j = 0; j < mainFoldLoopTimes; j++) {
278 AscendC::MicroAPI::RegTensor<float> reg0, reg1, reg0_1, reg1_1, reg2, reg2_1;281 AscendC::MicroAPI::RegTensor<float> reg0, reg1, reg0_1, reg1_1, reg2, reg2_1;
279 LoadTensorForDtypeTIn(foldX0A, reg0, pFull, i * outerLoopStride + j * innerLoopStride);282 LoadTensorForDtypeTIn(foldX0A, reg0, pFull, i * outerLoopStride + j * innerLoopStride);
@@ -286,8 +289,8 @@ __aicore__ inline void SoftmaxGradAR<T>::NormCompute(const int64_t aSize)
286 Mul(reg2_1, reg1, reg1_1, pFull);289 Mul(reg2_1, reg1, reg1_1, pFull);
287 290 
288 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2, reg2, reg2_1, pFull);291 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2, reg2, reg2_1, pFull);
289- ReduceSum(reg2, reg2, pFull);292+ Reduce<ReduceType::SUM>(reg2, reg2, pFull);
290- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, reg2, UReg, 1);293+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, reg2, UReg, 1);
291 }294 }
292 for (uint16_t j = 0; j < tailFoldLoopTimes; j++) {295 for (uint16_t j = 0; j < tailFoldLoopTimes; j++) {
293 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);296 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);
@@ -304,9 +307,9 @@ __aicore__ inline void SoftmaxGradAR<T>::NormCompute(const int64_t aSize)
304 Mul(reg2_1, reg1, reg1_1, pMask);307 Mul(reg2_1, reg1, reg1_1, pMask);
305 308 
306 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2_1, reg2, reg2_1, pMask);309 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2_1, reg2, reg2_1, pMask);
307- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg2, reg2_1, pMask);310+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg2, reg2_1, pMask);
308- ReduceSum(reg2, reg2, pFull);311+ Reduce<ReduceType::SUM>(reg2, reg2, pFull);
309- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, reg2, UReg, 1);312+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, reg2, UReg, 1);
310 }313 }
311 for (uint16_t j = 0; j < unFoldLoopTimes; j++) {314 for (uint16_t j = 0; j < unFoldLoopTimes; j++) {
312 AscendC::MicroAPI::RegTensor<float> reg0, reg1, reg0_1;315 AscendC::MicroAPI::RegTensor<float> reg0, reg1, reg0_1;
@@ -314,10 +317,10 @@ __aicore__ inline void SoftmaxGradAR<T>::NormCompute(const int64_t aSize)
314 LoadTensorForDtypeTIn(unFoldX1, reg0_1, pFull, i * outerLoopStride + j * innerLoopStride);317 LoadTensorForDtypeTIn(unFoldX1, reg0_1, pFull, i * outerLoopStride + j * innerLoopStride);
315 318 
316 Mul(reg1, reg0, reg0_1, pFull);319 Mul(reg1, reg0, reg0_1, pFull);
317- ReduceSum(reg1, reg1, pFull);320+ Reduce<ReduceType::SUM>(reg1, reg1, pFull);
318- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, reg1, UReg, 1);321+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, reg1, UReg, 1);
319 }322 }
320- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);323+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
321 }324 }
322 }325 }
323 NormComputePost(dstTensor, x0Tensor, x1Tensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride);326 NormComputePost(dstTensor, x0Tensor, x1Tensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride);
@@ -346,10 +349,10 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputePost(const LocalTensor<T>& d
346 uint16_t oriRAligned = tl_->rAligned;349 uint16_t oriRAligned = tl_->rAligned;
347 350 
348 if (rSize <= VL_FP32) {351 if (rSize <= VL_FP32) {
349- __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr();352+ __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr();
350- __local_mem__ T* x0 = (__local_mem__ T*)x0Tensor.GetPhyAddr();353+ __ubuf__ T* x0 = (__ubuf__ T*)x0Tensor.GetPhyAddr();
351- __local_mem__ T* x1 = (__local_mem__ T*)x1Tensor.GetPhyAddr();354+ __ubuf__ T* x1 = (__ubuf__ T*)x1Tensor.GetPhyAddr();
352- __local_mem__ float* sumTmp = (__local_mem__ float*)binAddTmpTensor.GetPhyAddr();355+ __ubuf__ float* sumTmp = (__ubuf__ float*)binAddTmpTensor.GetPhyAddr();
353 356 
354 __VEC_SCOPE__357 __VEC_SCOPE__
355 {358 {
@@ -360,8 +363,8 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputePost(const LocalTensor<T>& d
360 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();363 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
361 AscendC::MicroAPI::MaskReg maskOri;364 AscendC::MicroAPI::MaskReg maskOri;
362 for (uint16_t i = 0; i < loopTimes; i++) {365 for (uint16_t i = 0; i < loopTimes; i++) {
363- DataCopy(reg0, (__local_mem__ float*)sumTmp + i * static_cast<uint32_t>(stride));366+ LoadAlign(reg0, (__ubuf__ float*)sumTmp + i * static_cast<uint32_t>(stride));
364- ReduceSum(reg1, reg0, pMask);367+ Reduce<ReduceType::SUM>(reg1, reg0, pMask);
365 Duplicate(reg2, reg1, pFull);368 Duplicate(reg2, reg1, pFull);
366 369 
367 uint32_t sreg0 = static_cast<uint32_t>(oriR);370 uint32_t sreg0 = static_cast<uint32_t>(oriR);
@@ -378,12 +381,12 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputePost(const LocalTensor<T>& d
378 }381 }
379 }382 }
380 } else {383 } else {
381- __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr();384+ __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr();
382- __local_mem__ float* sumTmpA = (__local_mem__ float*)binAddTmpTensor.GetPhyAddr();385+ __ubuf__ float* sumTmpA = (__ubuf__ float*)binAddTmpTensor.GetPhyAddr();
383- __local_mem__ float* sumTmpB = (__local_mem__ float*)binAddTmpTensor.GetPhyAddr() + VL_FP32;386+ __ubuf__ float* sumTmpB = (__ubuf__ float*)binAddTmpTensor.GetPhyAddr() + VL_FP32;
384 387 
385- __local_mem__ T* x0 = (__local_mem__ T*)x0Tensor.GetPhyAddr();388+ __ubuf__ T* x0 = (__ubuf__ T*)x0Tensor.GetPhyAddr();
386- __local_mem__ T* x1 = (__local_mem__ T*)x1Tensor.GetPhyAddr();389+ __ubuf__ T* x1 = (__ubuf__ T*)x1Tensor.GetPhyAddr();
387 390 
388 __VEC_SCOPE__391 __VEC_SCOPE__
389 {392 {
@@ -394,11 +397,11 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputePost(const LocalTensor<T>& d
394 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();397 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
395 AscendC::MicroAPI::MaskReg maskOri;398 AscendC::MicroAPI::MaskReg maskOri;
396 for (uint16_t i = 0; i < loopTimes; i++) {399 for (uint16_t i = 0; i < loopTimes; i++) {
397- DataCopy(reg0, (__local_mem__ float*)sumTmpA + i * static_cast<uint32_t>(stride));400+ LoadAlign(reg0, (__ubuf__ float*)sumTmpA + i * static_cast<uint32_t>(stride));
398- DataCopy(reg1, (__local_mem__ float*)sumTmpB + i * static_cast<uint32_t>(stride));401+ LoadAlign(reg1, (__ubuf__ float*)sumTmpB + i * static_cast<uint32_t>(stride));
399 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg1, reg0, reg1, pMask);402 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg1, reg0, reg1, pMask);
400- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pMask);403+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pMask);
401- ReduceSum(reg2, reg0, pFull);404+ Reduce<ReduceType::SUM>(reg2, reg0, pFull);
402 Duplicate(reg2, reg2, pFull);405 Duplicate(reg2, reg2, pFull);
403 uint32_t sreg0 = static_cast<uint32_t>(oriR);406 uint32_t sreg0 = static_cast<uint32_t>(oriR);
404 for (uint16_t j = 0; j < rLoopCount; j++) {407 for (uint16_t j = 0; j < rLoopCount; j++) {
@@ -419,14 +422,14 @@ __aicore__ inline void SoftmaxGradAR<T>::NormComputePost(const LocalTensor<T>& d
419}422}
420 423 
421template <typename T>424template <typename T>
422-__aicore__ inline void SoftmaxGradAR<T>::LoadTensorForDtypeTIn(__local_mem__ T* src, RegTensor<float>& dst,425+__aicore__ inline void SoftmaxGradAR<T>::LoadTensorForDtypeTIn(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg,
423- MaskReg& preg, uint32_t offset)426+ uint32_t offset)
424{427{
425 if constexpr (IsSameType<T, float>::value) {428 if constexpr (IsSameType<T, float>::value) {
426- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);429+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
427 } else { // fp16、bf16430 } else { // fp16、bf16
428 RegTensor<T> xFp16;431 RegTensor<T> xFp16;
429- DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset));432+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset));
430 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);433 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);
431 }434 }
432}435}
@@ -451,16 +454,16 @@ __aicore__ inline void SoftmaxGradAR<T>::CopyInX(int64_t ubA, int64_t offset)
451}454}
452 455 
453template <typename T>456template <typename T>
454-__aicore__ inline void SoftmaxGradAR<T>::StoreTensorForDtypeTOut(__local_mem__ T* dst,457+__aicore__ inline void SoftmaxGradAR<T>::StoreTensorForDtypeTOut(__ubuf__ T* dst,
455 AscendC::MicroAPI::RegTensor<float>& src,458 AscendC::MicroAPI::RegTensor<float>& src,
456 AscendC::MicroAPI::MaskReg& preg, uint32_t offset)459 AscendC::MicroAPI::MaskReg& preg, uint32_t offset)
457{460{
458 if constexpr (IsSameType<T, float>::value) {461 if constexpr (IsSameType<T, float>::value) {
459- DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg);462+ StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg);
460 } else {463 } else {
461 AscendC::MicroAPI::RegTensor<T> xFp16;464 AscendC::MicroAPI::RegTensor<T> xFp16;
462 Cast<T, float, castTraitFp32ToFp16>(xFp16, src, preg);465 Cast<T, float, castTraitFp32ToFp16>(xFp16, src, preg);
463- DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg);466+ StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg);
464 }467 }
465}468}
466 469 
Mactivation/softmax_grad/op_kernel/arch35/softmax_grad_ar_recompute.h+29-27
@@ -24,6 +24,8 @@
24 24 
25namespace SoftmaxGradOps {25namespace SoftmaxGradOps {
26using namespace AscendC;26using namespace AscendC;
27+using AscendC::Reg::LoadAlign;
28+using AscendC::Reg::StoreAlign;
27 29 
28static constexpr int64_t AR_RECOMPUTE_SUM_BUFFER_BTYES = 32;30static constexpr int64_t AR_RECOMPUTE_SUM_BUFFER_BTYES = 32;
29static constexpr int64_t AR_RECOMPUTE_BINARY_CACHE_BTYES = 2048;31static constexpr int64_t AR_RECOMPUTE_BINARY_CACHE_BTYES = 2048;
@@ -42,12 +44,12 @@ private:
42 __aicore__ inline void CalcReduceSum(int64_t xDimOffset);44 __aicore__ inline void CalcReduceSum(int64_t xDimOffset);
43 __aicore__ inline void CalcOutVF(uint32_t ubFactor);45 __aicore__ inline void CalcOutVF(uint32_t ubFactor);
44 46 
45- __aicore__ inline void MainBlockVF(__local_mem__ float* dst, uint32_t ubFactor);47+ __aicore__ inline void MainBlockVF(__ubuf__ float* dst, uint32_t ubFactor);
46- __aicore__ inline void FoldBlockVF(__local_mem__ float* dst, uint32_t ubFactor);48+ __aicore__ inline void FoldBlockVF(__ubuf__ float* dst, uint32_t ubFactor);
47 49 
48- __aicore__ inline void LoadTensorForDtypeT(__local_mem__ T* src, RegTensor<float>& dst, MaskReg& pregMask,50+ __aicore__ inline void LoadTensorForDtypeT(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& pregMask,
49 uint32_t offset);51 uint32_t offset);
50- __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T* dst, AscendC::MicroAPI::RegTensor<float>& src,52+ __aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T* dst, AscendC::MicroAPI::RegTensor<float>& src,
51 AscendC::MicroAPI::MaskReg& preg, uint32_t offset);53 AscendC::MicroAPI::MaskReg& preg, uint32_t offset);
52 54 
53 __aicore__ inline void CopyInX(int64_t xGmOffset, uint32_t ubFactor);55 __aicore__ inline void CopyInX(int64_t xGmOffset, uint32_t ubFactor);
@@ -138,7 +140,7 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::CalcReduceSum(int64_t xDimOffs
138 LocalTensor<float> xSum = xSumBuffer_.Get<float>();140 LocalTensor<float> xSum = xSumBuffer_.Get<float>();
139 141 
140 LocalTensor<float> xTmp = yQueue_.AllocTensor<float>(); // 复用y做二分累加142 LocalTensor<float> xTmp = yQueue_.AllocTensor<float>(); // 复用y做二分累加
141- __local_mem__ float* xTmpLocal = (__local_mem__ float*)xTmp.GetPhyAddr();143+ __ubuf__ float* xTmpLocal = (__ubuf__ float*)xTmp.GetPhyAddr();
142 144 
143 // ub间累加fold折叠到main145 // ub间累加fold折叠到main
144 for (uint64_t basicBlockIdx = 0; basicBlockIdx < tl_->basicBlockLoop; basicBlockIdx++) {146 for (uint64_t basicBlockIdx = 0; basicBlockIdx < tl_->basicBlockLoop; basicBlockIdx++) {
@@ -178,13 +180,13 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::CalcReduceSum(int64_t xDimOffs
178 180 
179// cast + mul181// cast + mul
180template <typename T>182template <typename T>
181-__aicore__ inline void SoftmaxGradArRecompute<T>::MainBlockVF(__local_mem__ float* dst, uint32_t ubFactor)183+__aicore__ inline void SoftmaxGradArRecompute<T>::MainBlockVF(__ubuf__ float* dst, uint32_t ubFactor)
182{184{
183 LocalTensor<T> x0 = x0Queue_.DeQue<T>();185 LocalTensor<T> x0 = x0Queue_.DeQue<T>();
184 LocalTensor<T> x1 = x1Queue_.DeQue<T>();186 LocalTensor<T> x1 = x1Queue_.DeQue<T>();
185 187 
186- __local_mem__ T* x0Local = (__local_mem__ T*)x0.GetPhyAddr();188+ __ubuf__ T* x0Local = (__ubuf__ T*)x0.GetPhyAddr();
187- __local_mem__ T* x1Local = (__local_mem__ T*)x1.GetPhyAddr();189+ __ubuf__ T* x1Local = (__ubuf__ T*)x1.GetPhyAddr();
188 190 
189 __VEC_SCOPE__191 __VEC_SCOPE__
190 {192 {
@@ -201,7 +203,7 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::MainBlockVF(__local_mem__ floa
201 203 
202 Mul(reg0, reg0, reg1, pregMask);204 Mul(reg0, reg0, reg1, pregMask);
203 205 
204- AscendC::MicroAPI::DataCopy(dst + offset, reg0, pregMask);206+ AscendC::MicroAPI::StoreAlign(dst + offset, reg0, pregMask);
205 }207 }
206 }208 }
207 209 
@@ -210,13 +212,13 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::MainBlockVF(__local_mem__ floa
210}212}
211 213 
212template <typename T>214template <typename T>
213-__aicore__ inline void SoftmaxGradArRecompute<T>::FoldBlockVF(__local_mem__ float* dst, uint32_t ubFactor)215+__aicore__ inline void SoftmaxGradArRecompute<T>::FoldBlockVF(__ubuf__ float* dst, uint32_t ubFactor)
214{216{
215 LocalTensor<T> x0 = x0Queue_.DeQue<T>();217 LocalTensor<T> x0 = x0Queue_.DeQue<T>();
216 LocalTensor<T> x1 = x1Queue_.DeQue<T>();218 LocalTensor<T> x1 = x1Queue_.DeQue<T>();
217 219 
218- __local_mem__ T* x0Local = (__local_mem__ T*)x0.GetPhyAddr();220+ __ubuf__ T* x0Local = (__ubuf__ T*)x0.GetPhyAddr();
219- __local_mem__ T* x1Local = (__local_mem__ T*)x1.GetPhyAddr();221+ __ubuf__ T* x1Local = (__ubuf__ T*)x1.GetPhyAddr();
220 222 
221 __VEC_SCOPE__223 __VEC_SCOPE__
222 {224 {
@@ -236,12 +238,12 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::FoldBlockVF(__local_mem__ floa
236 238 
237 Mul(reg1, reg0, reg1, pregMask);239 Mul(reg1, reg0, reg1, pregMask);
238 240 
239- AscendC::MicroAPI::DataCopy(reg0, dst + offset);241+ AscendC::MicroAPI::LoadAlign(reg0, dst + offset);
240 242 
241 AscendC::MicroAPI::Add(reg1, reg0, reg1, pregMask);243 AscendC::MicroAPI::Add(reg1, reg0, reg1, pregMask);
242- AscendC::MicroAPI::Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pregMask);244+ AscendC::MicroAPI::Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pregMask);
243 245 
244- AscendC::MicroAPI::DataCopy(dst + offset, reg0, maskFull);246+ AscendC::MicroAPI::StoreAlign(dst + offset, reg0, maskFull);
245 }247 }
246 }248 }
247 249 
@@ -256,10 +258,10 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::CalcOutVF(uint32_t ubFactor)
256 LocalTensor<T> x1 = x1Queue_.DeQue<T>();258 LocalTensor<T> x1 = x1Queue_.DeQue<T>();
257 LocalTensor<T> y = yQueue_.AllocTensor<T>();259 LocalTensor<T> y = yQueue_.AllocTensor<T>();
258 260 
259- __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr();261+ __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr();
260- __local_mem__ T* x0Local = (__local_mem__ T*)x0.GetPhyAddr();262+ __ubuf__ T* x0Local = (__ubuf__ T*)x0.GetPhyAddr();
261- __local_mem__ T* x1Local = (__local_mem__ T*)x1.GetPhyAddr();263+ __ubuf__ T* x1Local = (__ubuf__ T*)x1.GetPhyAddr();
262- __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr();264+ __ubuf__ T* yLocal = (__ubuf__ T*)y.GetPhyAddr();
263 265 
264 __VEC_SCOPE__266 __VEC_SCOPE__
265 {267 {
@@ -269,7 +271,7 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::CalcOutVF(uint32_t ubFactor)
269 uint32_t sreg = ubFactor;271 uint32_t sreg = ubFactor;
270 uint16_t loopTimes = CeilDivision(ubFactor, VL_FP32);272 uint16_t loopTimes = CeilDivision(ubFactor, VL_FP32);
271 273 
272- AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, xSumLocal);274+ AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, xSumLocal);
273 275 
274 for (uint16_t j = 0; j < loopTimes; j++) {276 for (uint16_t j = 0; j < loopTimes; j++) {
275 pregMask = AscendC::MicroAPI::UpdateMask<float>(sreg);277 pregMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
@@ -291,30 +293,30 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::CalcOutVF(uint32_t ubFactor)
291}293}
292 294 
293template <typename T>295template <typename T>
294-__aicore__ inline void SoftmaxGradArRecompute<T>::LoadTensorForDtypeT(__local_mem__ T* src, RegTensor<float>& dst,296+__aicore__ inline void SoftmaxGradArRecompute<T>::LoadTensorForDtypeT(__ubuf__ T* src, RegTensor<float>& dst,
295 MaskReg& pregMask, uint32_t offset)297 MaskReg& pregMask, uint32_t offset)
296{298{
297 if constexpr (IsSameType<T, float>::value) {299 if constexpr (IsSameType<T, float>::value) {
298- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);300+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
299 } else { // fp16、bf16301 } else { // fp16、bf16
300 RegTensor<T> xFp16;302 RegTensor<T> xFp16;
301- DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset));303+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset));
302 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, pregMask);304 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, pregMask);
303 }305 }
304}306}
305 307 
306template <typename T>308template <typename T>
307-__aicore__ inline void SoftmaxGradArRecompute<T>::StoreTensorForDtypeTOut(__local_mem__ T* dst,309+__aicore__ inline void SoftmaxGradArRecompute<T>::StoreTensorForDtypeTOut(__ubuf__ T* dst,
308 AscendC::MicroAPI::RegTensor<float>& src,310 AscendC::MicroAPI::RegTensor<float>& src,
309 AscendC::MicroAPI::MaskReg& preg,311 AscendC::MicroAPI::MaskReg& preg,
310 uint32_t offset)312 uint32_t offset)
311{313{
312 if constexpr (IsSameType<T, float>::value) {314 if constexpr (IsSameType<T, float>::value) {
313- DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg);315+ StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg);
314 } else {316 } else {
315 AscendC::MicroAPI::RegTensor<T> xFp16;317 AscendC::MicroAPI::RegTensor<T> xFp16;
316 Cast<T, float, castTraitFp32ToFp16>(xFp16, src, preg);318 Cast<T, float, castTraitFp32ToFp16>(xFp16, src, preg);
317- DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg);319+ StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg);
318 }320 }
319}321}
320 322 
@@ -351,4 +353,4 @@ __aicore__ inline void SoftmaxGradArRecompute<T>::CopyInX(int64_t xGmOffset, uin
351 x1Queue_.EnQue(x1);353 x1Queue_.EnQue(x1);
352}354}
353} // namespace SoftmaxGradOps355} // namespace SoftmaxGradOps
354-#endif // SOFTMAX_GRAD_AR_RECOMPUTE_H356+#endif // SOFTMAX_GRAD_AR_RECOMPUTE_H
Mactivation/softmax_grad/op_kernel/arch35/softmax_grad_ar_small_r.h+29-27
@@ -30,6 +30,8 @@ using AscendC::MicroAPI::MaskMergeMode;
30using AscendC::MicroAPI::MaskReg;30using AscendC::MicroAPI::MaskReg;
31using AscendC::MicroAPI::RegTensor;31using AscendC::MicroAPI::RegTensor;
32using AscendC::MicroAPI::StoreDist;32using AscendC::MicroAPI::StoreDist;
33+using AscendC::Reg::LoadAlign;
34+using AscendC::Reg::StoreAlign;
33 35 
34template <typename T>36template <typename T>
35class SoftmaxGradARSmallR {37class SoftmaxGradARSmallR {
@@ -87,8 +89,8 @@ public:
87 LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>();89 LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>();
88 LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>();90 LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>();
89 91 
90- __local_mem__ T* x0Local = (__local_mem__ T*)x0Tensor.GetPhyAddr();92+ __ubuf__ T* x0Local = (__ubuf__ T*)x0Tensor.GetPhyAddr();
91- __local_mem__ T* x1Local = (__local_mem__ T*)x1Tensor.GetPhyAddr();93+ __ubuf__ T* x1Local = (__ubuf__ T*)x1Tensor.GetPhyAddr();
92 CalcReduceSum(x0Local, x1Local, curTileA0Len);94 CalcReduceSum(x0Local, x1Local, curTileA0Len);
93 CopyInAndTransPose(xOffsetPreLoad, nextTileA0Len, tilingData_->totalRLen);95 CopyInAndTransPose(xOffsetPreLoad, nextTileA0Len, tilingData_->totalRLen);
94 x1Queue_.FreeTensor(x1Tensor);96 x1Queue_.FreeTensor(x1Tensor);
@@ -104,8 +106,8 @@ public:
104 xOffset = curIdx * tilingData_->tileA0Len * tilingData_->totalRLen;106 xOffset = curIdx * tilingData_->tileA0Len * tilingData_->totalRLen;
105 LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>();107 LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>();
106 LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>();108 LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>();
107- __local_mem__ T* x0Local = (__local_mem__ T*)x0Tensor.GetPhyAddr();109+ __ubuf__ T* x0Local = (__ubuf__ T*)x0Tensor.GetPhyAddr();
108- __local_mem__ T* x1Local = (__local_mem__ T*)x1Tensor.GetPhyAddr();110+ __ubuf__ T* x1Local = (__ubuf__ T*)x1Tensor.GetPhyAddr();
109 CalcReduceSum(x0Local, x1Local, curTileA0Len);111 CalcReduceSum(x0Local, x1Local, curTileA0Len);
110 x1Queue_.FreeTensor(x1Tensor);112 x1Queue_.FreeTensor(x1Tensor);
111 CalcOutput(x0Local, curTileA0Len);113 CalcOutput(x0Local, curTileA0Len);
@@ -115,12 +117,11 @@ public:
115 }117 }
116 118 
117private:119private:
118- __aicore__ inline void CalcReduceSum(const __local_mem__ T* x0Local, const __local_mem__ T* x1Local,120+ __aicore__ inline void CalcReduceSum(const __ubuf__ T* x0Local, const __ubuf__ T* x1Local, uint32_t curTileA0Len)
119- uint32_t curTileA0Len)
120 {121 {
121- __local_mem__ float* tmpAddr = (__local_mem__ float*)tmpLocal_.GetPhyAddr();122+ __ubuf__ float* tmpAddr = (__ubuf__ float*)tmpLocal_.GetPhyAddr();
122- __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned]123+ __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned]
123- .GetPhyAddr();124+ .GetPhyAddr();
124 125 
125 uint32_t tileA0Len = tilingData_->tileA0Len;126 uint32_t tileA0Len = tilingData_->tileA0Len;
126 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);127 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);
@@ -141,8 +142,8 @@ private:
141 LoadTensorForDtypeT(x1Local, x1Reg, pregMask, xOffset);142 LoadTensorForDtypeT(x1Local, x1Reg, pregMask, xOffset);
142 143 
143 Mul(x0Reg, x0Reg, x1Reg, pregMask);144 Mul(x0Reg, x0Reg, x1Reg, pregMask);
144- DataCopy(tmpAddr + xOffset, x0Reg, pregMask);145+ StoreAlign(tmpAddr + xOffset, x0Reg, pregMask);
145- DataCopy(tmpAddr2 + xOffset, x0Reg, pregMask);146+ StoreAlign(tmpAddr2 + xOffset, x0Reg, pregMask);
146 }147 }
147 }148 }
148 }149 }
@@ -152,13 +153,13 @@ private:
152 AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(xSumTensor_, tmpLocal_, srcShape, false);153 AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(xSumTensor_, tmpLocal_, srcShape, false);
153 }154 }
154 155 
155- __aicore__ inline void CalcOutput(const __local_mem__ T* x0Local, uint32_t curTileA0Len)156+ __aicore__ inline void CalcOutput(const __ubuf__ T* x0Local, uint32_t curTileA0Len)
156 {157 {
157- __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr();158+ __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr();
158- __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned]159+ __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned]
159- .GetPhyAddr();160+ .GetPhyAddr();
160 tmpLocalTy_ = tmpLocal_.template ReinterpretCast<T>();161 tmpLocalTy_ = tmpLocal_.template ReinterpretCast<T>();
161- __local_mem__ T* tmpAddrTy = (__local_mem__ T*)tmpLocalTy_.GetPhyAddr();162+ __ubuf__ T* tmpAddrTy = (__ubuf__ T*)tmpLocalTy_.GetPhyAddr();
162 163 
163 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);164 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);
164 uint16_t loopA0Num = static_cast<uint16_t>(ops::CeilDiv(curTileA0Len, VL_FP32));165 uint16_t loopA0Num = static_cast<uint16_t>(ops::CeilDiv(curTileA0Len, VL_FP32));
@@ -174,21 +175,22 @@ private:
174 uint32_t tileA0LenLocal = tilingData_->tileA0Len;175 uint32_t tileA0LenLocal = tilingData_->tileA0Len;
175 for (uint16_t k = 0; k < loopA0Num; k++) {176 for (uint16_t k = 0; k < loopA0Num; k++) {
176 pregMask = UpdateMask<float>(sreg);177 pregMask = UpdateMask<float>(sreg);
177- DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32);178+ LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32);
178 for (uint16_t i = 0; i < curTileRLenVl; i++) {179 for (uint16_t i = 0; i < curTileRLenVl; i++) {
179 uint32_t xOffset = i * tileA0LenLocal + k * VL_FP32;180 uint32_t xOffset = i * tileA0LenLocal + k * VL_FP32;
180 LoadTensorForDtypeT(x0Local, x0Reg, pregMask, xOffset);181 LoadTensorForDtypeT(x0Local, x0Reg, pregMask, xOffset);
181 182 
182- DataCopy(x1Reg, tmpAddr2 + xOffset);183+ LoadAlign(x1Reg, tmpAddr2 + xOffset);
183 Neg(x0Reg, x0Reg, pregMask);184 Neg(x0Reg, x0Reg, pregMask);
184 MulAddDst(x1Reg, x0Reg, sumReg, pregMask);185 MulAddDst(x1Reg, x0Reg, sumReg, pregMask);
185 186 
186 if constexpr (xToFp32_) {187 if constexpr (xToFp32_) {
187- MicroAPI::DataCopy(tmpAddrTy + xOffset, x1Reg, pregMask);188+ MicroAPI::StoreAlign(tmpAddrTy + xOffset, x1Reg, pregMask);
188 } else { // fp16、bf16189 } else { // fp16、bf16
189 RegTensor<T> xFp16;190 RegTensor<T> xFp16;
190 MicroAPI::Cast<T, float, castTraitFp32ToFp16>(xFp16, x1Reg, pregMask);191 MicroAPI::Cast<T, float, castTraitFp32ToFp16>(xFp16, x1Reg, pregMask);
191- MicroAPI::DataCopy<T, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + xOffset, xFp16, pregMask);192+ MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + xOffset, xFp16,
193+ pregMask);
192 }194 }
193 }195 }
194 }196 }
@@ -260,29 +262,29 @@ private:
260 yQueue_.EnQue(yLocal_);262 yQueue_.EnQue(yLocal_);
261 }263 }
262 264 
263- __aicore__ inline void LoadTensorForDtypeT(const __local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg,265+ __aicore__ inline void LoadTensorForDtypeT(const __ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg,
264 uint32_t offset)266 uint32_t offset)
265 {267 {
266 if constexpr (xToFp32_) {268 if constexpr (xToFp32_) {
267- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);269+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
268 } else { // fp16、bf16270 } else { // fp16、bf16
269 RegTensor<T> xFp16;271 RegTensor<T> xFp16;
270- DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset));272+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset));
271 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);273 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);
272 }274 }
273 }275 }
274 276 
275 __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen)277 __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen)
276 {278 {
277- static constexpr MultiCopyConfig config = {false};279+ static constexpr NdDmaConfig config = {false};
278- MultiCopyLoopInfo<CONST_TWO> copyLoopInfo;280+ NdDmaLoopInfo<CONST_TWO> copyLoopInfo;
279 copyLoopInfo.loopSrcStride[0] = 1;281 copyLoopInfo.loopSrcStride[0] = 1;
280 copyLoopInfo.loopSrcStride[1] = totalRLen;282 copyLoopInfo.loopSrcStride[1] = totalRLen;
281 copyLoopInfo.loopDstStride[0] = tilingData_->tileA0Len;283 copyLoopInfo.loopDstStride[0] = tilingData_->tileA0Len;
282 copyLoopInfo.loopDstStride[1] = 1;284 copyLoopInfo.loopDstStride[1] = 1;
283 copyLoopInfo.loopSize[0] = totalRLen;285 copyLoopInfo.loopSize[0] = totalRLen;
284 copyLoopInfo.loopSize[1] = curTileA0Len;286 copyLoopInfo.loopSize[1] = curTileA0Len;
285- MultiCopyParams<T, CONST_TWO> params = {copyLoopInfo, 0};287+ NdDmaParams<T, CONST_TWO> params = {copyLoopInfo, 0};
286 288 
287 LocalTensor<T> x0Local_ = x0Queue_.AllocTensor<T>();289 LocalTensor<T> x0Local_ = x0Queue_.AllocTensor<T>();
288 DataCopy<T, CONST_TWO, config>(x0Local_, x0Gm_[xGmOffset], params);290 DataCopy<T, CONST_TWO, config>(x0Local_, x0Gm_[xGmOffset], params);
@@ -352,4 +354,4 @@ private:
352};354};
353} // namespace SoftmaxGradOps355} // namespace SoftmaxGradOps
354 356 
355-#endif357+#endif
Mactivation/softmax_grad/op_kernel/arch35/softmax_grad_ara_full_load.h+19-17
@@ -30,6 +30,8 @@ using AscendC::MicroAPI::MaskMergeMode;
30using AscendC::MicroAPI::MaskReg;30using AscendC::MicroAPI::MaskReg;
31using AscendC::MicroAPI::RegTensor;31using AscendC::MicroAPI::RegTensor;
32using AscendC::MicroAPI::StoreDist;32using AscendC::MicroAPI::StoreDist;
33+using AscendC::Reg::LoadAlign;
34+using AscendC::Reg::StoreAlign;
33 35 
34template <typename T>36template <typename T>
35class SoftmaxGradARA {37class SoftmaxGradARA {
@@ -88,8 +90,8 @@ public:
88 LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>();90 LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>();
89 LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>();91 LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>();
90 92 
91- __local_mem__ T* x0Local = (__local_mem__ T*)x0Tensor.GetPhyAddr();93+ __ubuf__ T* x0Local = (__ubuf__ T*)x0Tensor.GetPhyAddr();
92- __local_mem__ T* x1Local = (__local_mem__ T*)x1Tensor.GetPhyAddr();94+ __ubuf__ T* x1Local = (__ubuf__ T*)x1Tensor.GetPhyAddr();
93 95 
94 yMain_ = yQueue_.AllocTensor<float>();96 yMain_ = yQueue_.AllocTensor<float>();
95 97 
@@ -109,10 +111,10 @@ public:
109 }111 }
110 112 
111private:113private:
112- __aicore__ inline void CalcReduceSum(const __local_mem__ T* x0Local, const __local_mem__ T* x1Local,114+ __aicore__ inline void CalcReduceSum(const __ubuf__ T* x0Local, const __ubuf__ T* x1Local, uint32_t curTileA0Len,
113- uint32_t curTileA0Len, int64_t a0BlockOffset)115+ int64_t a0BlockOffset)
114 {116 {
115- __local_mem__ float* yLocal = (__local_mem__ float*)yMain_.GetPhyAddr() + a0BlockOffset;117+ __ubuf__ float* yLocal = (__ubuf__ float*)yMain_.GetPhyAddr() + a0BlockOffset;
116 118 
117 uint32_t tileA0Len = tilingData_->tileA0Len;119 uint32_t tileA0Len = tilingData_->tileA0Len;
118 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);120 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);
@@ -133,7 +135,7 @@ private:
133 LoadTensorForDtypeT(x1Local, x1Reg, pregMask, xOffset);135 LoadTensorForDtypeT(x1Local, x1Reg, pregMask, xOffset);
134 136 
135 Mul(x0Reg, x0Reg, x1Reg, pregMask);137 Mul(x0Reg, x0Reg, x1Reg, pregMask);
136- DataCopy(((__local_mem__ float*)yLocal) + xOffset, x0Reg, pregMask);138+ StoreAlign(((__ubuf__ float*)yLocal) + xOffset, x0Reg, pregMask);
137 }139 }
138 }140 }
139 }141 }
@@ -144,11 +146,11 @@ private:
144 false);146 false);
145 }147 }
146 148 
147- __aicore__ inline void CalcOutput(const __local_mem__ T* x0Local, const __local_mem__ T* x1Local,149+ __aicore__ inline void CalcOutput(const __ubuf__ T* x0Local, const __ubuf__ T* x1Local, uint32_t curTileA0Len,
148- uint32_t curTileA0Len, int64_t a0BlockOffset)150+ int64_t a0BlockOffset)
149 {151 {
150- __local_mem__ T* yLocal = (__local_mem__ T*)yMain_.GetPhyAddr() + a0BlockOffset;152+ __ubuf__ T* yLocal = (__ubuf__ T*)yMain_.GetPhyAddr() + a0BlockOffset;
151- __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr();153+ __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr();
152 154 
153 uint32_t tileA0Len = tilingData_->tileA0Len;155 uint32_t tileA0Len = tilingData_->tileA0Len;
154 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);156 uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen);
@@ -165,7 +167,7 @@ private:
165 167 
166 for (uint16_t k = 0; k < loopA0Num; k++) {168 for (uint16_t k = 0; k < loopA0Num; k++) {
167 pregMask = UpdateMask<float>(sreg);169 pregMask = UpdateMask<float>(sreg);
168- DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32);170+ LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32);
169 for (uint16_t i = 0; i < curTileRLenVl; i++) {171 for (uint16_t i = 0; i < curTileRLenVl; i++) {
170 uint32_t xOffset = i * tileA0Len + k * VL_FP32;172 uint32_t xOffset = i * tileA0Len + k * VL_FP32;
171 LoadTensorForDtypeT(x0Local, x0Reg, pregMask, xOffset);173 LoadTensorForDtypeT(x0Local, x0Reg, pregMask, xOffset);
@@ -176,25 +178,25 @@ private:
176 MulAddDst(x1Reg, x0Reg, sumReg, pregMask);178 MulAddDst(x1Reg, x0Reg, sumReg, pregMask);
177 179 
178 if constexpr (IsSameType<T, float>::value) {180 if constexpr (IsSameType<T, float>::value) {
179- DataCopy(((__local_mem__ float*)yLocal) + xOffset, x1Reg, pregMask);181+ StoreAlign(((__ubuf__ float*)yLocal) + xOffset, x1Reg, pregMask);
180 } else { // fp16、bf16182 } else { // fp16、bf16
181 RegTensor<T> xFp16;183 RegTensor<T> xFp16;
182 Cast<T, float, castTraitFp32ToFp16>(xFp16, x1Reg, pregMask);184 Cast<T, float, castTraitFp32ToFp16>(xFp16, x1Reg, pregMask);
183- DataCopy<T, StoreDist::DIST_PACK_B32>(((__local_mem__ T*)yLocal) + xOffset, xFp16, pregMask);185+ StoreAlign<T, StoreDist::DIST_PACK_B32>(((__ubuf__ T*)yLocal) + xOffset, xFp16, pregMask);
184 }186 }
185 }187 }
186 }188 }
187 }189 }
188 }190 }
189 191 
190- __aicore__ inline void LoadTensorForDtypeT(const __local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg,192+ __aicore__ inline void LoadTensorForDtypeT(const __ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg,
191 uint32_t offset)193 uint32_t offset)
192 {194 {
193 if constexpr (IsSameType<T, float>::value) {195 if constexpr (IsSameType<T, float>::value) {
194- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);196+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
195 } else { // fp16、bf16197 } else { // fp16、bf16
196 RegTensor<T> xFp16;198 RegTensor<T> xFp16;
197- DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset));199+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset));
198 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);200 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);
199 }201 }
200 }202 }
@@ -273,4 +275,4 @@ private:
273};275};
274} // namespace SoftmaxGradOps276} // namespace SoftmaxGradOps
275 277 
276-#endif278+#endif
Mactivation/softmax_grad/op_kernel/arch35/softmax_grad_ara_recompute.h+21-20
@@ -30,6 +30,8 @@ using AscendC::MicroAPI::MaskMergeMode;
30using AscendC::MicroAPI::MaskReg;30using AscendC::MicroAPI::MaskReg;
31using AscendC::MicroAPI::RegTensor;31using AscendC::MicroAPI::RegTensor;
32using AscendC::MicroAPI::StoreDist;32using AscendC::MicroAPI::StoreDist;
33+using AscendC::Reg::LoadAlign;
34+using AscendC::Reg::StoreAlign;
33 35 
34template <typename T>36template <typename T>
35class SoftmaxGradARARecompute : public SoftmaxGradOpsBase {37class SoftmaxGradARARecompute : public SoftmaxGradOpsBase {
@@ -147,9 +149,9 @@ private:
147 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);149 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);
148 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;150 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;
149 151 
150- __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr();152+ __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr();
151- __local_mem__ T* x0 = (__local_mem__ T*)x0Main_.GetPhyAddr();153+ __ubuf__ T* x0 = (__ubuf__ T*)x0Main_.GetPhyAddr();
152- __local_mem__ T* x1 = (__local_mem__ T*)x1Main_.GetPhyAddr();154+ __ubuf__ T* x1 = (__ubuf__ T*)x1Main_.GetPhyAddr();
153 155 
154 __VEC_SCOPE__156 __VEC_SCOPE__
155 {157 {
@@ -166,7 +168,7 @@ private:
166 LoadTensorForDtypeT(x0, x0Reg, pregMask, xOffset);168 LoadTensorForDtypeT(x0, x0Reg, pregMask, xOffset);
167 LoadTensorForDtypeT(x1, x1Reg, pregMask, xOffset);169 LoadTensorForDtypeT(x1, x1Reg, pregMask, xOffset);
168 Mul(x0Reg, x0Reg, x1Reg, pregMask);170 Mul(x0Reg, x0Reg, x1Reg, pregMask);
169- DataCopy((__local_mem__ float*)dst + xOffset, x0Reg, pregMask);171+ StoreAlign((__ubuf__ float*)dst + xOffset, x0Reg, pregMask);
170 }172 }
171 }173 }
172 }174 }
@@ -187,9 +189,9 @@ private:
187 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);189 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);
188 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;190 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;
189 191 
190- __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr();192+ __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr();
191- __local_mem__ T* x0 = (__local_mem__ T*)x0Fold.GetPhyAddr();193+ __ubuf__ T* x0 = (__ubuf__ T*)x0Fold.GetPhyAddr();
192- __local_mem__ T* x1 = (__local_mem__ T*)x1Fold.GetPhyAddr();194+ __ubuf__ T* x1 = (__ubuf__ T*)x1Fold.GetPhyAddr();
193 195 
194 __VEC_SCOPE__196 __VEC_SCOPE__
195 {197 {
@@ -206,9 +208,9 @@ private:
206 LoadTensorForDtypeT(x0, x0Reg, pregMask, i * outerLoopSrcStride + j * VL_FP32);208 LoadTensorForDtypeT(x0, x0Reg, pregMask, i * outerLoopSrcStride + j * VL_FP32);
207 LoadTensorForDtypeT(x1, x1Reg, pregMask, i * outerLoopSrcStride + j * VL_FP32);209 LoadTensorForDtypeT(x1, x1Reg, pregMask, i * outerLoopSrcStride + j * VL_FP32);
208 Mul(x0Reg, x0Reg, x1Reg, pregMask);210 Mul(x0Reg, x0Reg, x1Reg, pregMask);
209- DataCopy(x1Reg, (__local_mem__ float*)dst + xOffset);211+ LoadAlign(x1Reg, (__ubuf__ float*)dst + xOffset);
210 Add(x1Reg, x1Reg, x0Reg, pregMask);212 Add(x1Reg, x1Reg, x0Reg, pregMask);
211- DataCopy((__local_mem__ float*)dst + xOffset, x1Reg, pregMask);213+ StoreAlign((__ubuf__ float*)dst + xOffset, x1Reg, pregMask);
212 }214 }
213 }215 }
214 }216 }
@@ -248,14 +250,14 @@ private:
248 __aicore__ inline void CalcOutput(int64_t curTileRLen, uint32_t curTileA0Len, uint16_t loopA0Num)250 __aicore__ inline void CalcOutput(int64_t curTileRLen, uint32_t curTileA0Len, uint16_t loopA0Num)
249 {251 {
250 LocalTensor<T> x0 = x0Queue_.DeQue<T>();252 LocalTensor<T> x0 = x0Queue_.DeQue<T>();
251- __local_mem__ T* x0Local = (__local_mem__ T*)x0.GetPhyAddr();253+ __ubuf__ T* x0Local = (__ubuf__ T*)x0.GetPhyAddr();
252 LocalTensor<T> x1 = x1Queue_.DeQue<T>();254 LocalTensor<T> x1 = x1Queue_.DeQue<T>();
253- __local_mem__ T* x1Local = (__local_mem__ T*)x1.GetPhyAddr();255+ __ubuf__ T* x1Local = (__ubuf__ T*)x1.GetPhyAddr();
254 256 
255 LocalTensor<T> y = yQueue_.template AllocTensor<T>();257 LocalTensor<T> y = yQueue_.template AllocTensor<T>();
256- __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr();258+ __ubuf__ T* yLocal = (__ubuf__ T*)y.GetPhyAddr();
257 259 
258- __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr();260+ __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr();
259 261 
260 uint32_t tileA0Len = tilingData_->tileA0Len;262 uint32_t tileA0Len = tilingData_->tileA0Len;
261 uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen);263 uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen);
@@ -270,7 +272,7 @@ private:
270 272 
271 for (uint16_t k = 0; k < loopA0Num; k++) {273 for (uint16_t k = 0; k < loopA0Num; k++) {
272 pregMask = UpdateMask<float>(sreg);274 pregMask = UpdateMask<float>(sreg);
273- DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32);275+ LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32);
274 for (uint16_t i = 0; i < curTileRLenVl; i++) {276 for (uint16_t i = 0; i < curTileRLenVl; i++) {
275 uint32_t xOffset = i * tileA0Len + k * VL_FP32;277 uint32_t xOffset = i * tileA0Len + k * VL_FP32;
276 LoadTensorForDtypeT(x0Local, x0Reg, pregMask, xOffset);278 LoadTensorForDtypeT(x0Local, x0Reg, pregMask, xOffset);
@@ -282,11 +284,11 @@ private:
282 284 
283 // copy out285 // copy out
284 if constexpr (IsSameType<T, float>::value) {286 if constexpr (IsSameType<T, float>::value) {
285- DataCopy(((__local_mem__ float*)yLocal) + xOffset, x1Reg, pregMask);287+ StoreAlign(((__ubuf__ float*)yLocal) + xOffset, x1Reg, pregMask);
286 } else { // fp16、bf16288 } else { // fp16、bf16
287 RegTensor<T> xFp16;289 RegTensor<T> xFp16;
288 Cast<T, float, castTraitFp32ToFp16>(xFp16, x1Reg, pregMask);290 Cast<T, float, castTraitFp32ToFp16>(xFp16, x1Reg, pregMask);
289- DataCopy<T, StoreDist::DIST_PACK_B32>(((__local_mem__ T*)yLocal) + xOffset, xFp16, pregMask);291+ StoreAlign<T, StoreDist::DIST_PACK_B32>(((__ubuf__ T*)yLocal) + xOffset, xFp16, pregMask);
290 }292 }
291 }293 }
292 }294 }
@@ -298,14 +300,13 @@ private:
298 x1Queue_.FreeTensor<T>(x1);300 x1Queue_.FreeTensor<T>(x1);
299 }301 }
300 302 
301- __aicore__ inline void LoadTensorForDtypeT(__local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg,303+ __aicore__ inline void LoadTensorForDtypeT(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)
302- uint32_t offset)
303 {304 {
304 if constexpr (IsSameType<T, float>::value) {305 if constexpr (IsSameType<T, float>::value) {
305- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);306+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
306 } else { // fp16、bf16307 } else { // fp16、bf16
307 RegTensor<T> xFp16;308 RegTensor<T> xFp16;
308- DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset));309+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset));
309 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);310 Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg);
310 }311 }
311 }312 }
Mactivation/softmax_grad/op_kernel/arch35/softmax_grad_base.h+170-167
@@ -23,6 +23,10 @@
23 23 
24namespace SoftmaxGradOps {24namespace SoftmaxGradOps {
25using namespace AscendC;25using namespace AscendC;
26+using AscendC::Reg::LoadAlign;
27+using AscendC::Reg::Move;
28+using AscendC::Reg::Reduce;
29+using AscendC::Reg::StoreAlign;
26 30 
27constexpr static AscendC::MicroAPI::CastTrait castTraitFp16ToFp32 = {31constexpr static AscendC::MicroAPI::CastTrait castTraitFp16ToFp32 = {
28 AscendC::MicroAPI::RegLayout::ZERO,32 AscendC::MicroAPI::RegLayout::ZERO,
@@ -175,8 +179,8 @@ __aicore__ inline void SoftmaxGradOpsBase::CastToFp32From(const LocalTensor<floa
175 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {179 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
176 __VEC_SCOPE__180 __VEC_SCOPE__
177 {181 {
178- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();182+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
179- __local_mem__ T* src = (__local_mem__ T*)srcTensor.GetPhyAddr();183+ __ubuf__ T* src = (__ubuf__ T*)srcTensor.GetPhyAddr();
180 uint32_t count;184 uint32_t count;
181 AscendC::MicroAPI::RegTensor<float> fp32Reg;185 AscendC::MicroAPI::RegTensor<float> fp32Reg;
182 AscendC::MicroAPI::RegTensor<T> b16Reg;186 AscendC::MicroAPI::RegTensor<T> b16Reg;
@@ -185,10 +189,10 @@ __aicore__ inline void SoftmaxGradOpsBase::CastToFp32From(const LocalTensor<floa
185 count = static_cast<uint32_t>(colSize);189 count = static_cast<uint32_t>(colSize);
186 for (uint16_t j = 0; j < innerLoopTimes; ++j) {190 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
187 pMask = AscendC::MicroAPI::UpdateMask<float>(count);191 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
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 + j * innerLoopStride);193+ b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + j * innerLoopStride);
190 Cast<float, T, castTraitFp16ToFp32>(fp32Reg, b16Reg, pMask);194 Cast<float, T, castTraitFp16ToFp32>(fp32Reg, b16Reg, pMask);
191- DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask);195+ StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask);
192 }196 }
193 }197 }
194 }198 }
@@ -218,8 +222,8 @@ __aicore__ inline void SoftmaxGradOpsBase::CastFromFp32To(const LocalTensor<T>&
218 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {222 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
219 __VEC_SCOPE__223 __VEC_SCOPE__
220 {224 {
221- __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr();225+ __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr();
222- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();226+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
223 uint32_t count;227 uint32_t count;
224 AscendC::MicroAPI::RegTensor<float> fp32Reg;228 AscendC::MicroAPI::RegTensor<float> fp32Reg;
225 AscendC::MicroAPI::RegTensor<T> b16Reg;229 AscendC::MicroAPI::RegTensor<T> b16Reg;
@@ -228,10 +232,10 @@ __aicore__ inline void SoftmaxGradOpsBase::CastFromFp32To(const LocalTensor<T>&
228 count = static_cast<uint32_t>(colSize);232 count = static_cast<uint32_t>(colSize);
229 for (uint16_t j = 0; j < innerLoopTimes; ++j) {233 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
230 pMask = AscendC::MicroAPI::UpdateMask<float>(count);234 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
231- DataCopy(fp32Reg, (__local_mem__ float*)src + i * outerLoopSrcStride + j * innerLoopStride);235+ LoadAlign(fp32Reg, (__ubuf__ float*)src + i * outerLoopSrcStride + j * innerLoopStride);
232 Cast<T, float, castTraitFp32ToFp16>(b16Reg, fp32Reg, pMask);236 Cast<T, float, castTraitFp32ToFp16>(b16Reg, fp32Reg, pMask);
233- DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(237+ StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(
234- (__local_mem__ T*)dst + i * outerLoopDstStride + j * innerLoopStride, b16Reg, pMask);238+ (__ubuf__ T*)dst + i * outerLoopDstStride + j * innerLoopStride, b16Reg, pMask);
235 }239 }
236 }240 }
237 }241 }
@@ -317,19 +321,19 @@ __aicore__ inline void SoftmaxGradOpsBase::VectorAdd(const LocalTensor<float>& d
317 static_cast<int64_t>(platform::GetVRegSize()));321 static_cast<int64_t>(platform::GetVRegSize()));
318 __VEC_SCOPE__322 __VEC_SCOPE__
319 {323 {
320- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();324+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
321- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();325+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
322- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();326+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
323 uint32_t sreg = static_cast<uint32_t>(count);327 uint32_t sreg = static_cast<uint32_t>(count);
324 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;328 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
325 AscendC::MicroAPI::MaskReg pMask;329 AscendC::MicroAPI::MaskReg pMask;
326 for (uint16_t i = 0; i < loopTimes; ++i) {330 for (uint16_t i = 0; i < loopTimes; ++i) {
327 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);331 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
328- DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32);332+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32);
329- DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32);333+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32);
330 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);334 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
331- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);335+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
332- DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask);336+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask);
333 }337 }
334 }338 }
335}339}
@@ -347,20 +351,20 @@ __aicore__ inline void SoftmaxGradOpsBase::VectorAdd(const LocalTensor<float>& d
347 uint32_t innerLoopStride = stride;351 uint32_t innerLoopStride = stride;
348 __VEC_SCOPE__352 __VEC_SCOPE__
349 {353 {
350- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();354+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
351- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();355+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
352- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();356+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
353 uint32_t count = nSize;357 uint32_t count = nSize;
354 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;358 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
355 AscendC::MicroAPI::MaskReg pMask;359 AscendC::MicroAPI::MaskReg pMask;
356 for (uint16_t i = 0; i < outerLoopTimes; ++i) {360 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
357 pMask = AscendC::MicroAPI::UpdateMask<float>(count);361 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
358 for (uint16_t j = 0; j < innerLoopTimes; ++j) {362 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
359- DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride);363+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride);
360- DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + j * innerLoopStride);364+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + j * innerLoopStride);
361 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);365 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
362- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);366+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
363- DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask);367+ StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask);
364 }368 }
365 }369 }
366 }370 }
@@ -378,19 +382,19 @@ __aicore__ inline void SoftmaxGradOpsBase::VectorMul(const LocalTensor<float>& d
378 static_cast<int64_t>(platform::GetVRegSize()));382 static_cast<int64_t>(platform::GetVRegSize()));
379 __VEC_SCOPE__383 __VEC_SCOPE__
380 {384 {
381- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();385+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
382- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();386+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
383- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();387+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
384 uint32_t sreg = static_cast<uint32_t>(count);388 uint32_t sreg = static_cast<uint32_t>(count);
385 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;389 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
386 AscendC::MicroAPI::MaskReg pMask;390 AscendC::MicroAPI::MaskReg pMask;
387 391 
388 for (uint16_t i = 0; i < loopTimes; ++i) {392 for (uint16_t i = 0; i < loopTimes; ++i) {
389 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);393 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
390- DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32);394+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32);
391- DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32);395+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32);
392 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);396 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
393- DataCopy((__local_mem__ float*)dst + i * VL_FP32, cReg, pMask);397+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, cReg, pMask);
394 }398 }
395 }399 }
396}400}
@@ -414,19 +418,19 @@ __aicore__ inline void SoftmaxGradOpsBase::NlastBroadcastMul(const LocalTensor<f
414 uint32_t innerLoopStride = aSize;418 uint32_t innerLoopStride = aSize;
415 __VEC_SCOPE__419 __VEC_SCOPE__
416 {420 {
417- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();421+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
418- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();422+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
419- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();423+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
420 uint32_t count = static_cast<uint32_t>(aSize);424 uint32_t count = static_cast<uint32_t>(aSize);
421 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;425 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
422 AscendC::MicroAPI::MaskReg pMask;426 AscendC::MicroAPI::MaskReg pMask;
423 for (uint16_t i = 0; i < outerLoopTimes; ++i) {427 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
424 pMask = AscendC::MicroAPI::UpdateMask<float>(count);428 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
425- DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride);429+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride);
426 for (uint16_t j = 0; j < innerLoopTimes; ++j) {430 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
427- DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride);431+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride);
428 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);432 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
429- DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);433+ StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);
430 }434 }
431 }435 }
432 }436 }
@@ -451,40 +455,40 @@ __aicore__ inline void SoftmaxGradOpsBase::LastReduceSumSmallR(const LocalTensor
451 if (rSize <= VL_FP32) {455 if (rSize <= VL_FP32) {
452 __VEC_SCOPE__456 __VEC_SCOPE__
453 {457 {
454- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();458+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
455- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();459+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
456 uint32_t count = static_cast<uint32_t>(rSize);460 uint32_t count = static_cast<uint32_t>(rSize);
457 AscendC::MicroAPI::RegTensor<float> aReg, bReg;461 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
458 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);462 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
459- AscendC::MicroAPI::UnalignReg UReg;463+ AscendC::MicroAPI::UnalignRegForStore UReg;
460 for (uint16_t i = 0; i < loopTimes; ++i) {464 for (uint16_t i = 0; i < loopTimes; ++i) {
461- DataCopy(aReg, (__local_mem__ float*)src + i * stride);465+ LoadAlign(aReg, (__ubuf__ float*)src + i * stride);
462- ReduceSum(bReg, aReg, pMask);466+ Reduce<ReduceType::SUM>(bReg, aReg, pMask);
463- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);467+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
464 }468 }
465- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);469+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
466 }470 }
467 } else {471 } else {
468 __VEC_SCOPE__472 __VEC_SCOPE__
469 {473 {
470- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();474+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
471- __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr();475+ __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr();
472- __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32;476+ __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32;
473 uint32_t count = static_cast<uint32_t>(rSize - VL_FP32);477 uint32_t count = static_cast<uint32_t>(rSize - VL_FP32);
474 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;478 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
475- AscendC::MicroAPI::UnalignReg UReg;479+ AscendC::MicroAPI::UnalignRegForStore UReg;
476 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);480 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
477 AscendC::MicroAPI::MaskReg481 AscendC::MicroAPI::MaskReg
478 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();482 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
479 for (uint16_t i = 0; i < loopTimes; ++i) {483 for (uint16_t i = 0; i < loopTimes; ++i) {
480- DataCopy(aReg, (__local_mem__ float*)src0 + i * stride);484+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * stride);
481- DataCopy(bReg, (__local_mem__ float*)src1 + i * stride);485+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * stride);
482 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);486 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
483- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);487+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
484- ReduceSum(bReg, aReg, pFull);488+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
485- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);489+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
486 }490 }
487- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);491+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
488 }492 }
489 }493 }
490}494}
@@ -529,43 +533,43 @@ __aicore__ inline void SoftmaxGradOpsBase::LastReduceSum(const LocalTensor<float
529 533 
530 __VEC_SCOPE__534 __VEC_SCOPE__
531 {535 {
532- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr();536+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr();
533- __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr();537+ __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr();
534- __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset;538+ __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset;
535- __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset;539+ __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset;
536- __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset;540+ __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset;
537- __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset;541+ __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset;
538 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();542 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
539- AscendC::MicroAPI::UnalignReg UReg;543+ AscendC::MicroAPI::UnalignRegForStore UReg;
540 544 
541 for (uint16_t i = 0; i < outerLoopTimes; ++i) {545 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
542- dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;546+ dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;
543 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {547 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {
544 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg;548 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg;
545- DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride);549+ LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride);
546- DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride);550+ LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride);
547 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull);551 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull);
548- ReduceSum(dReg, cReg, pFull);552+ Reduce<ReduceType::SUM>(dReg, cReg, pFull);
549- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1);553+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1);
550 }554 }
551 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {555 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {
552 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);556 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);
553 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;557 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
554 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);558 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
555- DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride);559+ LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride);
556- DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride);560+ LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride);
557 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);561 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
558- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);562+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
559- ReduceSum(bReg, aReg, pFull);563+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
560- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);564+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
561 }565 }
562 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {566 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {
563 AscendC::MicroAPI::RegTensor<float> aReg, bReg;567 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
564- DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride);568+ LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride);
565- ReduceSum(bReg, aReg, pFull);569+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
566- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);570+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
567 }571 }
568- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);572+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
569 }573 }
570 }574 }
571 LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride);575 LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride);
@@ -573,13 +577,13 @@ __aicore__ inline void SoftmaxGradOpsBase::LastReduceSum(const LocalTensor<float
573 577 
574template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1>578template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1>
575struct NlastDichotomyAdd {579struct NlastDichotomyAdd {
576- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,580+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
577- __local_mem__ float*& srcA, __local_mem__ float*& srcB,581+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
578- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)582+ uint32_t stride)
579 {583 {
580 AscendC::MicroAPI::RegTensor<float> aReg, bReg;584 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
581- __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO;585+ __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO;
582- __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO;586+ __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO;
583 if constexpr (TailCount <= 0) {587 if constexpr (TailCount <= 0) {
584 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,588 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,
585 stride * CONST_TWO);589 stride * CONST_TWO);
@@ -587,13 +591,13 @@ struct NlastDichotomyAdd {
587 }591 }
588 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);592 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
589 }593 }
590- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,594+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
591- __local_mem__ float*& srcA, __local_mem__ float*& srcB,595+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
592- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset)596+ uint32_t stride, uint32_t offset)
593 {597 {
594 AscendC::MicroAPI::RegTensor<float> aReg, bReg;598 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
595- __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO;599+ __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO;
596- __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO;600+ __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO;
597 if constexpr (TailCount <= 0) {601 if constexpr (TailCount <= 0) {
598 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,602 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,
599 stride * CONST_TWO, offset);603 stride * CONST_TWO, offset);
@@ -611,49 +615,49 @@ struct NlastDichotomyAdd {
611 615 
612template <int32_t TailCount, int32_t Index, int32_t Depth>616template <int32_t TailCount, int32_t Index, int32_t Depth>
613struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> {617struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> {
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- DataCopy(aReg, (__local_mem__ float*)srcA);623+ LoadAlign(aReg, (__ubuf__ float*)srcA);
620- DataCopy(bReg, (__local_mem__ float*)srcB);624+ LoadAlign(bReg, (__ubuf__ float*)srcB);
621 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);625 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
622 }626 }
623- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,627+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
624- __local_mem__ float*& srcA, __local_mem__ float*& srcB,628+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
625- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset)629+ uint32_t stride, uint32_t offset)
626 {630 {
627 if constexpr (TailCount <= 0) {631 if constexpr (TailCount <= 0) {
628 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;632 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
629- DataCopy(aReg, (__local_mem__ float*)srcA);633+ LoadAlign(aReg, (__ubuf__ float*)srcA);
630- DataCopy(bReg, (__local_mem__ float*)srcA + offset);634+ LoadAlign(bReg, (__ubuf__ float*)srcA + offset);
631 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);635 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
632- DataCopy(bReg, (__local_mem__ float*)srcB);636+ LoadAlign(bReg, (__ubuf__ float*)srcB);
633- DataCopy(cReg, (__local_mem__ float*)srcB + offset);637+ LoadAlign(cReg, (__ubuf__ float*)srcB + offset);
634 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);638 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);
635 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);639 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
636 } else {640 } else {
637 if constexpr (Index + Depth < TailCount) {641 if constexpr (Index + Depth < TailCount) {
638 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;642 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
639- DataCopy(aReg, (__local_mem__ float*)srcA);643+ LoadAlign(aReg, (__ubuf__ float*)srcA);
640- DataCopy(bReg, (__local_mem__ float*)srcA + offset);644+ LoadAlign(bReg, (__ubuf__ float*)srcA + offset);
641 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);645 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
642- DataCopy(bReg, (__local_mem__ float*)srcB);646+ LoadAlign(bReg, (__ubuf__ float*)srcB);
643- DataCopy(cReg, (__local_mem__ float*)srcB + offset);647+ LoadAlign(cReg, (__ubuf__ float*)srcB + offset);
644 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);648 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);
645 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);649 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
646 } else if constexpr (Index < TailCount) {650 } else if constexpr (Index < TailCount) {
647 AscendC::MicroAPI::RegTensor<float> aReg, bReg;651 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
648- DataCopy(aReg, (__local_mem__ float*)srcA);652+ LoadAlign(aReg, (__ubuf__ float*)srcA);
649- DataCopy(bReg, (__local_mem__ float*)srcA + offset);653+ LoadAlign(bReg, (__ubuf__ float*)srcA + offset);
650 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);654 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
651- DataCopy(bReg, (__local_mem__ float*)srcB);655+ LoadAlign(bReg, (__ubuf__ float*)srcB);
652 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);656 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
653 } else {657 } else {
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*)srcB);660+ LoadAlign(bReg, (__ubuf__ float*)srcB);
657 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);661 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
658 }662 }
659 }663 }
@@ -662,25 +666,25 @@ struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> {
662 666 
663template <>667template <>
664struct NlastDichotomyAdd<CONST_TWO> {668struct NlastDichotomyAdd<CONST_TWO> {
665- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,669+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
666- __local_mem__ float*& srcA, __local_mem__ float*& srcB,670+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
667- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)671+ uint32_t stride)
668 {672 {
669 AscendC::MicroAPI::RegTensor<float> aReg, bReg;673 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
670- DataCopy(aReg, (__local_mem__ float*)srcA);674+ LoadAlign(aReg, (__ubuf__ float*)srcA);
671- DataCopy(bReg, (__local_mem__ float*)srcB);675+ LoadAlign(bReg, (__ubuf__ float*)srcB);
672 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);676 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
673 }677 }
674- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,678+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
675- __local_mem__ float*& srcA, __local_mem__ float*& srcB,679+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
676- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset)680+ uint32_t stride, uint32_t offset)
677 {681 {
678 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;682 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
679- DataCopy(aReg, (__local_mem__ float*)srcA);683+ LoadAlign(aReg, (__ubuf__ float*)srcA);
680- DataCopy(bReg, (__local_mem__ float*)srcA + offset);684+ LoadAlign(bReg, (__ubuf__ float*)srcA + offset);
681 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);685 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
682- DataCopy(bReg, (__local_mem__ float*)srcB);686+ LoadAlign(bReg, (__ubuf__ float*)srcB);
683- DataCopy(cReg, (__local_mem__ float*)srcB + offset);687+ LoadAlign(cReg, (__ubuf__ float*)srcB + offset);
684 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);688 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);
685 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);689 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
686 }690 }
@@ -688,11 +692,11 @@ struct NlastDichotomyAdd<CONST_TWO> {
688 692 
689template <>693template <>
690struct NlastDichotomyAdd<1> {694struct NlastDichotomyAdd<1> {
691- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,695+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
692- __local_mem__ float*& srcA, __local_mem__ float*& srcB,696+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
693- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)697+ uint32_t stride)
694 {698 {
695- DataCopy(acc, (__local_mem__ float*)srcA);699+ LoadAlign(acc, (__ubuf__ float*)srcA);
696 }700 }
697};701};
698 702 
@@ -707,32 +711,32 @@ __aicore__ inline void SoftmaxGradOpsBase::NlastReduceSumSmallR(const LocalTenso
707 if constexpr (RSize == 1) {711 if constexpr (RSize == 1) {
708 __VEC_SCOPE__712 __VEC_SCOPE__
709 {713 {
710- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();714+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
711- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();715+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
712 uint32_t count = static_cast<uint32_t>(aSize);716 uint32_t count = static_cast<uint32_t>(aSize);
713 AscendC::MicroAPI::RegTensor<float> aReg;717 AscendC::MicroAPI::RegTensor<float> aReg;
714 AscendC::MicroAPI::MaskReg pMask;718 AscendC::MicroAPI::MaskReg pMask;
715 for (uint16_t i = 0; i < loopTimes; ++i) {719 for (uint16_t i = 0; i < loopTimes; ++i) {
716 pMask = AscendC::MicroAPI::UpdateMask<float>(count);720 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
717- DataCopy(aReg, (__local_mem__ float*)src + i * VL_FP32);721+ LoadAlign(aReg, (__ubuf__ float*)src + i * VL_FP32);
718- DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask);722+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask);
719 }723 }
720 }724 }
721 } else {725 } else {
722 __VEC_SCOPE__726 __VEC_SCOPE__
723 {727 {
724- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();728+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
725- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr();729+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr();
726- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride;730+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride;
727 uint32_t count = static_cast<uint32_t>(aSize);731 uint32_t count = static_cast<uint32_t>(aSize);
728 AscendC::MicroAPI::RegTensor<float> aReg;732 AscendC::MicroAPI::RegTensor<float> aReg;
729 AscendC::MicroAPI::MaskReg pMask;733 AscendC::MicroAPI::MaskReg pMask;
730 for (uint16_t i = 0; i < loopTimes; ++i) {734 for (uint16_t i = 0; i < loopTimes; ++i) {
731- __local_mem__ float* curSrcA = srcA + i * VL_FP32;735+ __ubuf__ float* curSrcA = srcA + i * VL_FP32;
732- __local_mem__ float* curSrcB = srcB + i * VL_FP32;736+ __ubuf__ float* curSrcB = srcB + i * VL_FP32;
733 pMask = AscendC::MicroAPI::UpdateMask<float>(count);737 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
734 NlastDichotomyAdd<RSize>::LoadAndAccumulate(aReg, curSrcA, curSrcB, pMask, stride);738 NlastDichotomyAdd<RSize>::LoadAndAccumulate(aReg, curSrcA, curSrcB, pMask, stride);
735- DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask);739+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask);
736 }740 }
737 }741 }
738 }742 }
@@ -805,35 +809,34 @@ __aicore__ inline void SoftmaxGradOpsBase::NlastReduceSumLargeR(const LocalTenso
805 for (uint16_t i = 0; i < outerLoopTimes; ++i) {809 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
806 pMask = plt_b32(count, POST_UPDATE);810 pMask = plt_b32(count, POST_UPDATE);
807 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {811 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {
808- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() +812+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopStride +
809- i * outerLoopStride + j * innerLoopStride;813+ j * innerLoopStride;
810- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride +814+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride + j * srcStride;
811- j * srcStride;815+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride + i * outerLoopStride +
812- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride +816+ j * srcStride;
813- i * outerLoopStride + j * srcStride;
814 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, foldOffset);817 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, foldOffset);
815- DataCopy((__local_mem__ float*)dst, aReg, pMask);818+ StoreAlign((__ubuf__ float*)dst, aReg, pMask);
816 }819 }
817 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {820 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {
818- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() +821+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopStride +
819- i * outerLoopStride + mainFoldLoopTimes * innerLoopStride;822+ mainFoldLoopTimes * innerLoopStride;
820- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride +823+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride +
821- mainFoldLoopTimes * srcStride;824+ mainFoldLoopTimes * srcStride;
822- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride +825+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride + i * outerLoopStride +
823- i * outerLoopStride + mainFoldLoopTimes * srcStride;826+ mainFoldLoopTimes * srcStride;
824 NlastDichotomyAdd<COMPRESSION, TailCount>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride,827 NlastDichotomyAdd<COMPRESSION, TailCount>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride,
825 foldOffset);828 foldOffset);
826- DataCopy((__local_mem__ float*)dst, aReg, pMask);829+ StoreAlign((__ubuf__ float*)dst, aReg, pMask);
827 }830 }
828 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {831 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {
829- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + unFoldDstOffset +832+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + unFoldDstOffset +
830- i * outerLoopStride + j * innerLoopStride;833+ i * outerLoopStride + j * innerLoopStride;
831- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset +834+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + i * outerLoopStride +
832- i * outerLoopStride + j * srcStride;835+ j * srcStride;
833- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + stride +836+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + stride +
834- i * outerLoopStride + j * srcStride;837+ i * outerLoopStride + j * srcStride;
835 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride);838 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride);
836- DataCopy((__local_mem__ float*)dst, aReg, pMask);839+ StoreAlign((__ubuf__ float*)dst, aReg, pMask);
837 }840 }
838 }841 }
839 }842 }
@@ -894,20 +897,20 @@ __aicore__ inline void SoftmaxGradOpsBase::UpdateCache(const LocalTensor<float>&
894 uint32_t innerLoopStride = stride;897 uint32_t innerLoopStride = stride;
895 __VEC_SCOPE__898 __VEC_SCOPE__
896 {899 {
897- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();900+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
898- __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride;901+ __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride;
899- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();902+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
900 uint32_t sreg = static_cast<uint32_t>(count);903 uint32_t sreg = static_cast<uint32_t>(count);
901 AscendC::MicroAPI::RegTensor<float> aReg, bReg;904 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
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 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);907 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
905- DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride);908+ LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride);
906 for (uint16_t j = 0; j < innerLoopTimes; ++j) {909 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
907- DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride);910+ LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride);
908 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);911 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
909 }912 }
910- DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask);913+ StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask);
911 }914 }
912 }915 }
913}916}
@@ -926,27 +929,27 @@ __aicore__ inline void SoftmaxGradOpsBase::Normalize(const LocalTensor<float>& d
926 uint32_t innerLoopStride = VL_FP32;929 uint32_t innerLoopStride = VL_FP32;
927 __VEC_SCOPE__930 __VEC_SCOPE__
928 {931 {
929- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();932+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
930- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();933+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
931- __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr();934+ __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr();
932- __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr();935+ __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr();
933 uint32_t count;936 uint32_t count;
934 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;937 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
935 AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg;938 AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg;
936 AscendC::MicroAPI::MaskReg pMask;939 AscendC::MicroAPI::MaskReg pMask;
937 for (uint16_t i = 0; i < outerLoopTimes; ++i) {940 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
938 count = static_cast<uint32_t>(colSize);941 count = static_cast<uint32_t>(colSize);
939- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i);942+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i);
940- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i);943+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i);
941 for (uint16_t j = 0; j < innerLoopTimes; ++j) {944 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
942 pMask = AscendC::MicroAPI::UpdateMask<float>(count);945 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
943- DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + j * innerLoopStride);946+ LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + j * innerLoopStride);
944 Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask);947 Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask);
945 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask);948 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask);
946- DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);949+ StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);
947 }950 }
948 }951 }
949 }952 }
950}953}
951} // namespace SoftmaxGradOps954} // namespace SoftmaxGradOps
952-#endif955+#endif
Mactivation/softmax_v2/op_kernel/arch35/softmax_v2_ar_full_load.h+57-54
@@ -27,6 +27,10 @@
27 27 
28namespace SoftmaxV2Ops {28namespace SoftmaxV2Ops {
29using namespace AscendC;29using namespace AscendC;
30+using AscendC::Reg::LoadAlign;
31+using AscendC::Reg::Move;
32+using AscendC::Reg::Reduce;
33+using AscendC::Reg::StoreAlign;
30 34 
31constexpr static uint32_t DOUBLE_BUFFER = 2;35constexpr static uint32_t DOUBLE_BUFFER = 2;
32constexpr static uint32_t BLOCK_SIZE = 32; // 32B36constexpr static uint32_t BLOCK_SIZE = 32; // 32B
@@ -42,8 +46,7 @@ public:
42private:46private:
43 __aicore__ inline void ProcessUB(int64_t ubA, int64_t aOffset);47 __aicore__ inline void ProcessUB(int64_t ubA, int64_t aOffset);
44 48 
45- __aicore__ inline void FirstNormCompute(int64_t ubA, __local_mem__ T_in* xInAddr,49+ __aicore__ inline void FirstNormCompute(int64_t ubA, __ubuf__ T_in* xInAddr, __ubuf__ float* xTmpLocalAddr);
46- __local_mem__ float* xTmpLocalAddr);
47 __aicore__ inline void SecondNormCompute(const LocalTensor<T_out>& dstTensor, const LocalTensor<float>& srcTensor,50 __aicore__ inline void SecondNormCompute(const LocalTensor<T_out>& dstTensor, const LocalTensor<float>& srcTensor,
48 const LocalTensor<float>& reduceSumTempTensor, const int64_t aSize,51 const LocalTensor<float>& reduceSumTempTensor, const int64_t aSize,
49 const int64_t rSize, const int64_t stride);52 const int64_t rSize, const int64_t stride);
@@ -54,9 +57,9 @@ private:
54 57 
55 __aicore__ inline void CopyInX(const LocalTensor<T_in>& xInUb, int64_t ubA, int64_t offset);58 __aicore__ inline void CopyInX(const LocalTensor<T_in>& xInUb, int64_t ubA, int64_t offset);
56 __aicore__ inline void CopyOutY(const LocalTensor<T_out>& yOutUb, int64_t ubA, int64_t offset);59 __aicore__ inline void CopyOutY(const LocalTensor<T_out>& yOutUb, int64_t ubA, int64_t offset);
57- __aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_in* src, AscendC::MicroAPI::RegTensor<float>& dst,60+ __aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_in* src, AscendC::MicroAPI::RegTensor<float>& dst,
58 AscendC::MicroAPI::MaskReg& preg, uint32_t offset);61 AscendC::MicroAPI::MaskReg& preg, uint32_t offset);
59- __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_out* dst, AscendC::MicroAPI::RegTensor<float>& src,62+ __aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T_out* dst, AscendC::MicroAPI::RegTensor<float>& src,
60 AscendC::MicroAPI::MaskReg& preg, uint32_t offset);63 AscendC::MicroAPI::MaskReg& preg, uint32_t offset);
61 64 
62private:65private:
@@ -123,9 +126,9 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::ProcessUB(int64_t ubA, int64_t
123 LocalTensor<float> binaryTmpLocalTensor = tmpLocalTensor[tl_->ubFactor * tl_->rAligned];126 LocalTensor<float> binaryTmpLocalTensor = tmpLocalTensor[tl_->ubFactor * tl_->rAligned];
124 LocalTensor<float> xTmpLocalTensor = tmpLocalTensor[0];127 LocalTensor<float> xTmpLocalTensor = tmpLocalTensor[0];
125 128 
126- __local_mem__ T_in* xInUbAddr = (__local_mem__ T_in*)xInUb.GetPhyAddr();129+ __ubuf__ T_in* xInUbAddr = (__ubuf__ T_in*)xInUb.GetPhyAddr();
127- __local_mem__ float* xTmpLocalAddr = (__local_mem__ float*)xTmpLocalTensor.GetPhyAddr();130+ __ubuf__ float* xTmpLocalAddr = (__ubuf__ float*)xTmpLocalTensor.GetPhyAddr();
128- __local_mem__ float* binaryTmpLocalAddr = (__local_mem__ float*)binaryTmpLocalTensor.GetPhyAddr();131+ __ubuf__ float* binaryTmpLocalAddr = (__ubuf__ float*)binaryTmpLocalTensor.GetPhyAddr();
129 132 
130 FirstNormCompute(ubA, xInUbAddr, xTmpLocalAddr);133 FirstNormCompute(ubA, xInUbAddr, xTmpLocalAddr);
131 xQueue_.FreeTensor<T_in>(xInUb);134 xQueue_.FreeTensor<T_in>(xInUb);
@@ -139,8 +142,8 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::ProcessUB(int64_t ubA, int64_t
139}142}
140 143 
141template <typename T_in, typename T_out>144template <typename T_in, typename T_out>
142-__aicore__ inline void SoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, __local_mem__ T_in* xInAddr,145+__aicore__ inline void SoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, __ubuf__ T_in* xInAddr,
143- __local_mem__ float* xTmpLocalAddr)146+ __ubuf__ float* xTmpLocalAddr)
144{147{
145 int64_t rAligned = tl_->rAligned;148 int64_t rAligned = tl_->rAligned;
146 int64_t ubActualR = tl_->r;149 int64_t ubActualR = tl_->r;
@@ -173,7 +176,7 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, _
173 AscendC::MicroAPI::Duplicate(vreg0, static_cast<float>(-INFINITY), maskAll);176 AscendC::MicroAPI::Duplicate(vreg0, static_cast<float>(-INFINITY), maskAll);
174 LoadTensorForDtypeTIn(xInAddr, vreg1, maskTail, tailAddrPtr);177 LoadTensorForDtypeTIn(xInAddr, vreg1, maskTail, tailAddrPtr);
175 AscendC::MicroAPI::Max(vreg1, vreg0, vreg1, maskTail);178 AscendC::MicroAPI::Max(vreg1, vreg0, vreg1, maskTail);
176- AscendC::MicroAPI::Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vreg0, vreg1, maskTail);179+ AscendC::MicroAPI::Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vreg0, vreg1, maskTail);
177 180 
178 uint32_t sreg2 = static_cast<uint32_t>(ubActualR - tailUbBlockSize);181 uint32_t sreg2 = static_cast<uint32_t>(ubActualR - tailUbBlockSize);
179 for (uint16_t i = 0; i < rLoopCountTmp; i++) {182 for (uint16_t i = 0; i < rLoopCountTmp; i++) {
@@ -182,7 +185,7 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, _
182 LoadTensorForDtypeTIn(xInAddr, vreg1, maskMax, addrPtr);185 LoadTensorForDtypeTIn(xInAddr, vreg1, maskMax, addrPtr);
183 AscendC::MicroAPI::Max(vreg0, vreg0, vreg1, maskMax);186 AscendC::MicroAPI::Max(vreg0, vreg0, vreg1, maskMax);
184 }187 }
185- AscendC::MicroAPI::ReduceMax(vreg2, vreg0, maskAll);188+ AscendC::MicroAPI::Reduce<ReduceType::MAX>(vreg2, vreg0, maskAll);
186 AscendC::MicroAPI::Duplicate(vreg3, vreg2, maskAll);189 AscendC::MicroAPI::Duplicate(vreg3, vreg2, maskAll);
187 190 
188 // 求sub和exp191 // 求sub和exp
@@ -193,7 +196,7 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, _
193 LoadTensorForDtypeTIn(xInAddr, vreg1, maskSub, addrPtr);196 LoadTensorForDtypeTIn(xInAddr, vreg1, maskSub, addrPtr);
194 AscendC::MicroAPI::Sub(vreg4, vreg1, vreg3, maskSub);197 AscendC::MicroAPI::Sub(vreg4, vreg1, vreg3, maskSub);
195 AscendC::MicroAPI::Exp(vreg5, vreg4, maskSub);198 AscendC::MicroAPI::Exp(vreg5, vreg4, maskSub);
196- AscendC::MicroAPI::DataCopy(((__local_mem__ float*)xTmpLocalAddr + addrPtr), vreg5, maskSub);199+ AscendC::MicroAPI::StoreAlign(((__ubuf__ float*)xTmpLocalAddr + addrPtr), vreg5, maskSub);
197 }200 }
198 }201 }
199 }202 }
@@ -222,9 +225,9 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::SecondNormComputePost(const Loc
222 uint16_t oriRAligned = tl_->rAligned;225 uint16_t oriRAligned = tl_->rAligned;
223 226 
224 if (rSize <= VL_FP32) {227 if (rSize <= VL_FP32) {
225- __local_mem__ T_out* dst = (__local_mem__ T_out*)dstTensor.GetPhyAddr();228+ __ubuf__ T_out* dst = (__ubuf__ T_out*)dstTensor.GetPhyAddr();
226- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();229+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
227- __local_mem__ float* oriSrc = (__local_mem__ float*)oriSrcTensor.GetPhyAddr();230+ __ubuf__ float* oriSrc = (__ubuf__ float*)oriSrcTensor.GetPhyAddr();
228 231 
229 __VEC_SCOPE__232 __VEC_SCOPE__
230 {233 {
@@ -235,24 +238,24 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::SecondNormComputePost(const Loc
235 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();238 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
236 AscendC::MicroAPI::MaskReg maskOri;239 AscendC::MicroAPI::MaskReg maskOri;
237 for (uint16_t i = 0; i < loopTimes; ++i) {240 for (uint16_t i = 0; i < loopTimes; ++i) {
238- DataCopy(aReg, (__local_mem__ float*)src + i * static_cast<uint32_t>(stride));241+ LoadAlign(aReg, (__ubuf__ float*)src + i * static_cast<uint32_t>(stride));
239- ReduceSum(bReg, aReg, pMask);242+ Reduce<ReduceType::SUM>(bReg, aReg, pMask);
240 Duplicate(cReg, bReg, pFull);243 Duplicate(cReg, bReg, pFull);
241 uint32_t sreg0 = static_cast<uint32_t>(oriR);244 uint32_t sreg0 = static_cast<uint32_t>(oriR);
242 for (uint16_t j = 0; j < rLoopCount; ++j) {245 for (uint16_t j = 0; j < rLoopCount; ++j) {
243 maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0);246 maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0);
244 uint32_t addrPtr = j * VL_FP32 + i * oriRAligned;247 uint32_t addrPtr = j * VL_FP32 + i * oriRAligned;
245- DataCopy(dReg, ((__local_mem__ float*)oriSrc + addrPtr));248+ LoadAlign(dReg, ((__ubuf__ float*)oriSrc + addrPtr));
246 Div(dReg, dReg, cReg, maskOri);249 Div(dReg, dReg, cReg, maskOri);
247 StoreTensorForDtypeTOut(dst, dReg, maskOri, addrPtr);250 StoreTensorForDtypeTOut(dst, dReg, maskOri, addrPtr);
248 }251 }
249 }252 }
250 }253 }
251 } else {254 } else {
252- __local_mem__ T_out* dst = (__local_mem__ T_out*)dstTensor.GetPhyAddr();255+ __ubuf__ T_out* dst = (__ubuf__ T_out*)dstTensor.GetPhyAddr();
253- __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr();256+ __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr();
254- __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32;257+ __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32;
255- __local_mem__ float* oriSrc = (__local_mem__ float*)oriSrcTensor.GetPhyAddr();258+ __ubuf__ float* oriSrc = (__ubuf__ float*)oriSrcTensor.GetPhyAddr();
256 259 
257 __VEC_SCOPE__260 __VEC_SCOPE__
258 {261 {
@@ -263,17 +266,17 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::SecondNormComputePost(const Loc
263 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();266 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
264 AscendC::MicroAPI::MaskReg maskOri;267 AscendC::MicroAPI::MaskReg maskOri;
265 for (uint16_t i = 0; i < loopTimes; ++i) {268 for (uint16_t i = 0; i < loopTimes; ++i) {
266- DataCopy(aReg, (__local_mem__ float*)src0 + i * static_cast<uint32_t>(stride));269+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * static_cast<uint32_t>(stride));
267- DataCopy(bReg, (__local_mem__ float*)src1 + i * static_cast<uint32_t>(stride));270+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * static_cast<uint32_t>(stride));
268 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);271 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
269- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);272+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
270- ReduceSum(bReg, aReg, pFull);273+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
271 Duplicate(dReg, bReg, pFull);274 Duplicate(dReg, bReg, pFull);
272 uint32_t sreg0 = static_cast<uint32_t>(oriR);275 uint32_t sreg0 = static_cast<uint32_t>(oriR);
273 for (uint16_t j = 0; j < rLoopCount; ++j) {276 for (uint16_t j = 0; j < rLoopCount; ++j) {
274 maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0);277 maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0);
275 uint32_t addrPtr = j * VL_FP32 + i * oriRAligned;278 uint32_t addrPtr = j * VL_FP32 + i * oriRAligned;
276- DataCopy(eReg, ((__local_mem__ float*)oriSrc + addrPtr));279+ LoadAlign(eReg, ((__ubuf__ float*)oriSrc + addrPtr));
277 Div(eReg, eReg, dReg, maskOri);280 Div(eReg, eReg, dReg, maskOri);
278 StoreTensorForDtypeTOut(dst, eReg, maskOri, addrPtr);281 StoreTensorForDtypeTOut(dst, eReg, maskOri, addrPtr);
279 }282 }
@@ -321,78 +324,78 @@ __aicore__ inline void SoftmaxV2AR<T_in, T_out>::SecondNormCompute(const LocalTe
321 int64_t tailSrcBOffset = floorVLCount * VL_FP32;324 int64_t tailSrcBOffset = floorVLCount * VL_FP32;
322 int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32;325 int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32;
323 326 
324- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr();327+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr();
325- __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr();328+ __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr();
326- __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset;329+ __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset;
327- __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset;330+ __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset;
328- __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset;331+ __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset;
329- __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset;332+ __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset;
330 333 
331 __VEC_SCOPE__334 __VEC_SCOPE__
332 {335 {
333 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();336 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
334- AscendC::MicroAPI::UnalignReg UReg;337+ AscendC::MicroAPI::UnalignRegForStore UReg;
335 338 
336 for (uint16_t i = 0; i < outerLoopTimes; ++i) {339 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
337- dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;340+ dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;
338 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {341 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {
339 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg;342 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg;
340- DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride);343+ LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride);
341- DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride);344+ LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride);
342 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull);345 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull);
343- ReduceSum(dReg, cReg, pFull);346+ Reduce<ReduceType::SUM>(dReg, cReg, pFull);
344- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1);347+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1);
345 }348 }
346 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {349 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {
347 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);350 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);
348 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;351 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
349 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);352 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
350- DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride);353+ LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride);
351- DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride);354+ LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride);
352 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);355 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
353- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);356+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
354- ReduceSum(bReg, aReg, pFull);357+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
355- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);358+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
356 }359 }
357 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {360 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {
358 AscendC::MicroAPI::RegTensor<float> aReg, bReg;361 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
359- DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride);362+ LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride);
360- ReduceSum(bReg, aReg, pFull);363+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
361- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);364+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
362 }365 }
363- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);366+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
364 }367 }
365 }368 }
366 SecondNormComputePost(dstTensor, reduceSumTempTensor, srcTensor, aSize, foldPoint, outerLoopDstStride);369 SecondNormComputePost(dstTensor, reduceSumTempTensor, srcTensor, aSize, foldPoint, outerLoopDstStride);
367}370}
368 371 
369template <typename T_in, typename T_out>372template <typename T_in, typename T_out>
370-__aicore__ inline void SoftmaxV2AR<T_in, T_out>::LoadTensorForDtypeTIn(__local_mem__ T_in* src,373+__aicore__ inline void SoftmaxV2AR<T_in, T_out>::LoadTensorForDtypeTIn(__ubuf__ T_in* src,
371 AscendC::MicroAPI::RegTensor<float>& dst,374 AscendC::MicroAPI::RegTensor<float>& dst,
372 AscendC::MicroAPI::MaskReg& preg,375 AscendC::MicroAPI::MaskReg& preg,
373 uint32_t offset)376 uint32_t offset)
374{377{
375 if constexpr (IsSameType<T_in, float>::value) {378 if constexpr (IsSameType<T_in, float>::value) {
376- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(dst, src + offset);379+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(dst, src + offset);
377 } else {380 } else {
378 AscendC::MicroAPI::RegTensor<T_in> xFp16;381 AscendC::MicroAPI::RegTensor<T_in> xFp16;
379- DataCopy<T_in, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, src + offset);382+ LoadAlign<T_in, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, src + offset);
380 Cast<float, T_in, castTraitFp16ToFp32>(dst, xFp16, preg);383 Cast<float, T_in, castTraitFp16ToFp32>(dst, xFp16, preg);
381 }384 }
382}385}
383 386 
384template <typename T_in, typename T_out>387template <typename T_in, typename T_out>
385-__aicore__ inline void SoftmaxV2AR<T_in, T_out>::StoreTensorForDtypeTOut(__local_mem__ T_out* dst,388+__aicore__ inline void SoftmaxV2AR<T_in, T_out>::StoreTensorForDtypeTOut(__ubuf__ T_out* dst,
386 AscendC::MicroAPI::RegTensor<float>& src,389 AscendC::MicroAPI::RegTensor<float>& src,
387 AscendC::MicroAPI::MaskReg& preg,390 AscendC::MicroAPI::MaskReg& preg,
388 uint32_t offset)391 uint32_t offset)
389{392{
390 if constexpr (IsSameType<T_out, float>::value) {393 if constexpr (IsSameType<T_out, float>::value) {
391- DataCopy<T_out, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg);394+ StoreAlign<T_out, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg);
392 } else {395 } else {
393 AscendC::MicroAPI::RegTensor<T_out> xFp16;396 AscendC::MicroAPI::RegTensor<T_out> xFp16;
394 Cast<T_out, float, castTraitFp32ToFp16>(xFp16, src, preg);397 Cast<T_out, float, castTraitFp32ToFp16>(xFp16, src, preg);
395- DataCopy<T_out, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg);398+ StoreAlign<T_out, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg);
396 }399 }
397}400}
398 401 
Mactivation/softmax_v2/op_kernel/arch35/softmax_v2_ar_recompute.h+57-63
@@ -41,15 +41,14 @@ public:
41 __aicore__ inline void Process();41 __aicore__ inline void Process();
42 42 
43private:43private:
44- __aicore__ inline void CalculateMaxVF(__local_mem__ float*& xMaxPtr, __local_mem__ Tx*& xPtr, uint32_t aSize,44+ __aicore__ inline void CalculateMaxVF(__ubuf__ float*& xMaxPtr, __ubuf__ Tx*& xPtr, uint32_t aSize,
45 uint32_t ubFactor);45 uint32_t ubFactor);
46- __aicore__ inline void CalculateOutVF(__local_mem__ Ty*& yPtr, __local_mem__ Tx*& xPtr,46+ __aicore__ inline void CalculateOutVF(__ubuf__ Ty*& yPtr, __ubuf__ Tx*& xPtr, __ubuf__ float*& xMaxPtr,
47- __local_mem__ float*& xMaxPtr, __local_mem__ float*& xSumPtr, uint32_t a,47+ __ubuf__ float*& xSumPtr, uint32_t a, uint32_t ubFactor);
48- uint32_t ubFactor);48+ __aicore__ inline void MainBlockCastSubExpVF(__ubuf__ float*& xFp32Ptr, __ubuf__ Tx*& xPtr,
49- __aicore__ inline void MainBlockCastSubExpVF(__local_mem__ float*& xFp32Ptr, __local_mem__ Tx*& xPtr,49+ __ubuf__ float*& xMaxPtr, uint32_t a, uint32_t ubFactor);
50- __local_mem__ float*& xMaxPtr, uint32_t a, uint32_t ubFactor);50+ __aicore__ inline void FoldBlockCastSubExpVF(__ubuf__ float*& dstPtr, __ubuf__ Tx*& xPtr, __ubuf__ float*& xMaxPtr,
51- __aicore__ inline void FoldBlockCastSubExpVF(__local_mem__ float*& dstPtr, __local_mem__ Tx*& xPtr,51+ uint32_t a, uint32_t ubFactor);
52- __local_mem__ float*& xMaxPtr, uint32_t a, uint32_t ubFactor);
53 __aicore__ inline int64_t GetCacheId(const int64_t idx);52 __aicore__ inline int64_t GetCacheId(const int64_t idx);
54 __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,53 __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,
55 const int64_t cacheId, const int64_t stride, const int64_t count);54 const int64_t cacheId, const int64_t stride, const int64_t count);
@@ -133,7 +132,7 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
133 x1DataCopyExtParams.srcStride = 0;132 x1DataCopyExtParams.srcStride = 0;
134 x1DataCopyExtParams.dstStride = 0;133 x1DataCopyExtParams.dstStride = 0;
135 134 
136- __local_mem__ float* xMaxPtr = (__local_mem__ float*)xMaxLocal.GetPhyAddr();135+ __ubuf__ float* xMaxPtr = (__ubuf__ float*)xMaxLocal.GetPhyAddr();
137 // step 1. 对R循环，求整行R的最大值136 // step 1. 对R循环，求整行R的最大值
138 for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) {137 for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) {
139 int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; // 每个UB循环的偏移量138 int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; // 每个UB循环的偏移量
@@ -148,14 +147,14 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
148 xQueue_.EnQue<Tx>(xLocal);147 xQueue_.EnQue<Tx>(xLocal);
149 xLocal = xQueue_.DeQue<Tx>();148 xLocal = xQueue_.DeQue<Tx>();
150 149 
151- __local_mem__ Tx* xPtr = (__local_mem__ Tx*)xLocal.GetPhyAddr();150+ __ubuf__ Tx* xPtr = (__ubuf__ Tx*)xLocal.GetPhyAddr();
152 CalculateMaxVF(xMaxPtr, xPtr, A_IN_IN, ubFactor);151 CalculateMaxVF(xMaxPtr, xPtr, A_IN_IN, ubFactor);
153 xQueue_.FreeTensor(xLocal);152 xQueue_.FreeTensor(xLocal);
154 }153 }
155 154 
156 // step 2. UB间二分累加：计算每行的Σe^(x - max)155 // step 2. UB间二分累加：计算每行的Σe^(x - max)
157 LocalTensor<float> xTmpLocal = xTmpBuffer.Get<float>();156 LocalTensor<float> xTmpLocal = xTmpBuffer.Get<float>();
158- __local_mem__ float* xTmpFp32Ptr = (__local_mem__ float*)xTmpLocal.GetPhyAddr();157+ __ubuf__ float* xTmpFp32Ptr = (__ubuf__ float*)xTmpLocal.GetPhyAddr();
159 158 
160 x1DataCopyExtParams.blockLen = tl_->ubFactor * sizeof(Tx);159 x1DataCopyExtParams.blockLen = tl_->ubFactor * sizeof(Tx);
161 160 
@@ -176,14 +175,14 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
176 xQueue_.EnQue<Tx>(x1Local);175 xQueue_.EnQue<Tx>(x1Local);
177 x1Local = xQueue_.DeQue<Tx>();176 x1Local = xQueue_.DeQue<Tx>();
178 177 
179- __local_mem__ Tx* x1Ptr = (__local_mem__ Tx*)x1Local.GetPhyAddr();178+ __ubuf__ Tx* x1Ptr = (__ubuf__ Tx*)x1Local.GetPhyAddr();
180 MainBlockCastSubExpVF(xTmpFp32Ptr, x1Ptr, xMaxPtr, A_IN_IN, tl_->ubFactor);179 MainBlockCastSubExpVF(xTmpFp32Ptr, x1Ptr, xMaxPtr, A_IN_IN, tl_->ubFactor);
181 xQueue_.FreeTensor(x1Local);180 xQueue_.FreeTensor(x1Local);
182 181 
183 // 折叠部分：X2折叠到X1上182 // 折叠部分：X2折叠到X1上
184 if (basicBlockIdx < tl_->mainFoldCount) {183 if (basicBlockIdx < tl_->mainFoldCount) {
185 LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>();184 LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>();
186- __local_mem__ Tx* x2Ptr = (__local_mem__ Tx*)x2Local.GetPhyAddr();185+ __ubuf__ Tx* x2Ptr = (__ubuf__ Tx*)x2Local.GetPhyAddr();
187 DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams);186 DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams);
188 xQueue_.EnQue<Tx>(x2Local);187 xQueue_.EnQue<Tx>(x2Local);
189 x2Local = xQueue_.DeQue<Tx>();188 x2Local = xQueue_.DeQue<Tx>();
@@ -192,7 +191,7 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
192 xQueue_.FreeTensor(x2Local);191 xQueue_.FreeTensor(x2Local);
193 } else if ((basicBlockIdx == tl_->mainFoldCount) && (tl_->ubFactorTail > 0)) {192 } else if ((basicBlockIdx == tl_->mainFoldCount) && (tl_->ubFactorTail > 0)) {
194 LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>();193 LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>();
195- __local_mem__ Tx* x2Ptr = (__local_mem__ Tx*)x2Local.GetPhyAddr();194+ __ubuf__ Tx* x2Ptr = (__ubuf__ Tx*)x2Local.GetPhyAddr();
196 x2DataCopyExtParams.blockLen = tl_->ubFactorTail * sizeof(Tx); // 这里的x2为尾块195 x2DataCopyExtParams.blockLen = tl_->ubFactorTail * sizeof(Tx); // 这里的x2为尾块
197 DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams);196 DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams);
198 xQueue_.EnQue<Tx>(x2Local);197 xQueue_.EnQue<Tx>(x2Local);
@@ -213,7 +212,7 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
213 // R很小，不需要做UB间二分累加212 // R很小，不需要做UB间二分累加
214 if (tl_->basicBlockLoop == 0) {213 if (tl_->basicBlockLoop == 0) {
215 LocalTensor<Tx> x1Local = xQueue_.AllocTensor<Tx>();214 LocalTensor<Tx> x1Local = xQueue_.AllocTensor<Tx>();
216- __local_mem__ Tx* x1Ptr = (__local_mem__ Tx*)x1Local.GetPhyAddr();215+ __ubuf__ Tx* x1Ptr = (__ubuf__ Tx*)x1Local.GetPhyAddr();
217 DataCopyPad(x1Local[0], xGm_[xDimOffset], x1DataCopyExtParams, padExtParams);216 DataCopyPad(x1Local[0], xGm_[xDimOffset], x1DataCopyExtParams, padExtParams);
218 xQueue_.EnQue<Tx>(x1Local);217 xQueue_.EnQue<Tx>(x1Local);
219 x1Local = xQueue_.DeQue<Tx>();218 x1Local = xQueue_.DeQue<Tx>();
@@ -231,7 +230,7 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
231 yDataCopyExtParams.srcStride = 0;230 yDataCopyExtParams.srcStride = 0;
232 yDataCopyExtParams.dstStride = 0;231 yDataCopyExtParams.dstStride = 0;
233 232 
234- __local_mem__ float* xSumPtr = (__local_mem__ float*)totalSumLocal_.GetPhyAddr();233+ __ubuf__ float* xSumPtr = (__ubuf__ float*)totalSumLocal_.GetPhyAddr();
235 // step 3. 遍历UB块，计算除法234 // step 3. 遍历UB块，计算除法
236 for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) {235 for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) {
237 int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx;236 int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx;
@@ -242,8 +241,8 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
242 241 
243 LocalTensor<Tx> xLocal = xQueue_.AllocTensor<Tx>();242 LocalTensor<Tx> xLocal = xQueue_.AllocTensor<Tx>();
244 LocalTensor<Ty> yLocal = yQueue_.AllocTensor<Ty>();243 LocalTensor<Ty> yLocal = yQueue_.AllocTensor<Ty>();
245- __local_mem__ Tx* xPtr = (__local_mem__ Tx*)xLocal.GetPhyAddr();244+ __ubuf__ Tx* xPtr = (__ubuf__ Tx*)xLocal.GetPhyAddr();
246- __local_mem__ Ty* yPtr = (__local_mem__ Ty*)yLocal.GetPhyAddr();245+ __ubuf__ Ty* yPtr = (__ubuf__ Ty*)yLocal.GetPhyAddr();
247 246 
248 x1DataCopyExtParams.blockLen = ubFactor * sizeof(Tx);247 x1DataCopyExtParams.blockLen = ubFactor * sizeof(Tx);
249 DataCopyPad(xLocal[0], xGm_[xUbOffset], x1DataCopyExtParams, padExtParams);248 DataCopyPad(xLocal[0], xGm_[xUbOffset], x1DataCopyExtParams, padExtParams);
@@ -263,9 +262,8 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::Process()
263}262}
264 263 
265template <typename Tx, typename Ty>264template <typename Tx, typename Ty>
266-__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateMaxVF(__local_mem__ float*& xMaxPtr,265+__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateMaxVF(__ubuf__ float*& xMaxPtr, __ubuf__ Tx*& xPtr,
267- __local_mem__ Tx*& xPtr, uint32_t aSize,266+ uint32_t aSize, uint32_t ubFactor)
268- uint32_t ubFactor)
269{267{
270 __VEC_SCOPE__268 __VEC_SCOPE__
271 {269 {
@@ -289,38 +287,37 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateMaxVF(__local_mem_
289 AscendC::MicroAPI::Duplicate(maxReg, CONST_FP32_MIN);287 AscendC::MicroAPI::Duplicate(maxReg, CONST_FP32_MIN);
290 288 
291 if constexpr (xToFp32_) {289 if constexpr (xToFp32_) {
292- AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr);290+ AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr);
293 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskTail);291 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskTail);
294 } else {292 } else {
295- AscendC::MicroAPI::DataCopy(vreg1, xAddr);293+ AscendC::MicroAPI::LoadAlign(vreg1, xAddr);
296 }294 }
297 AscendC::MicroAPI::Max(vreg1, maxReg, vreg1, maskTail);295 AscendC::MicroAPI::Max(vreg1, maxReg, vreg1, maskTail);
298- AscendC::MicroAPI::Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(maxReg, vreg1, maskTail);296+ AscendC::MicroAPI::Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(maxReg, vreg1, maskTail);
299 297 
300 // 整块处理298 // 整块处理
301 for (uint16_t j = 0; j < repeatTimesTmp; j++) {299 for (uint16_t j = 0; j < repeatTimesTmp; j++) {
302 auto xAddr = xPtr + j * VL_FP32;300 auto xAddr = xPtr + j * VL_FP32;
303 if constexpr (xToFp32_) {301 if constexpr (xToFp32_) {
304- AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr);302+ AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr);
305 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskFull);303 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskFull);
306 } else {304 } else {
307- AscendC::MicroAPI::DataCopy(vreg1, xAddr);305+ AscendC::MicroAPI::LoadAlign(vreg1, xAddr);
308 }306 }
309 AscendC::MicroAPI::Max(maxReg, maxReg, vreg1, maskFull);307 AscendC::MicroAPI::Max(maxReg, maxReg, vreg1, maskFull);
310 }308 }
311- AscendC::MicroAPI::DataCopy(vreg2, xMaxPtr);309+ AscendC::MicroAPI::LoadAlign(vreg2, xMaxPtr);
312 310 
313- AscendC::MicroAPI::ReduceMax(maxReg, maxReg, maskFull);311+ AscendC::MicroAPI::Reduce<ReduceType::MAX>(maxReg, maxReg, maskFull);
314 AscendC::MicroAPI::Max(maxReg, maxReg, vreg2, maskOne);312 AscendC::MicroAPI::Max(maxReg, maxReg, vreg2, maskOne);
315- AscendC::MicroAPI::DataCopy(xMaxPtr, maxReg, maskOne);313+ AscendC::MicroAPI::StoreAlign(xMaxPtr, maxReg, maskOne);
316 }314 }
317}315}
318 316 
319template <typename Tx, typename Ty>317template <typename Tx, typename Ty>
320-__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_mem__ Ty*& yPtr, __local_mem__ Tx*& xPtr,318+__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__ubuf__ Ty*& yPtr, __ubuf__ Tx*& xPtr,
321- __local_mem__ float*& xMaxPtr,319+ __ubuf__ float*& xMaxPtr, __ubuf__ float*& xSumPtr,
322- __local_mem__ float*& xSumPtr, uint32_t a,320+ uint32_t a, uint32_t ubFactor)
323- uint32_t ubFactor)
324{321{
325 __VEC_SCOPE__322 __VEC_SCOPE__
326 {323 {
@@ -332,8 +329,8 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_mem_
332 uint32_t width = ubFactor;329 uint32_t width = ubFactor;
333 uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32);330 uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32);
334 331 
335- AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr);332+ AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr);
336- AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, xSumPtr);333+ AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, xSumPtr);
337 334 
338 for (uint16_t j = 0; j < repeatTimes; j++) {335 for (uint16_t j = 0; j < repeatTimes; j++) {
339 mask = AscendC::MicroAPI::UpdateMask<float>(width);336 mask = AscendC::MicroAPI::UpdateMask<float>(width);
@@ -341,10 +338,10 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_mem_
341 auto yAddr = yPtr + j * VL_FP32;338 auto yAddr = yPtr + j * VL_FP32;
342 339 
343 if constexpr (xToFp32_) {340 if constexpr (xToFp32_) {
344- AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr);341+ AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr);
345 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask);342 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask);
346 } else {343 } else {
347- AscendC::MicroAPI::DataCopy(vreg1, xAddr);344+ AscendC::MicroAPI::LoadAlign(vreg1, xAddr);
348 }345 }
349 346 
350 AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask);347 AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask);
@@ -352,20 +349,19 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_mem_
352 AscendC::MicroAPI::Div(vreg3, vreg2, sumReg, mask);349 AscendC::MicroAPI::Div(vreg3, vreg2, sumReg, mask);
353 350 
354 if constexpr (yToFp32_) {351 if constexpr (yToFp32_) {
355- AscendC::MicroAPI::DataCopy(yAddr, vreg3, mask);352+ AscendC::MicroAPI::StoreAlign(yAddr, vreg3, mask);
356 } else {353 } else {
357 AscendC::MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(vreg4, vreg3, mask);354 AscendC::MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(vreg4, vreg3, mask);
358- AscendC::MicroAPI::DataCopy<Ty, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(yAddr, vreg4, mask);355+ AscendC::MicroAPI::StoreAlign<Ty, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(yAddr, vreg4, mask);
359 }356 }
360 }357 }
361 }358 }
362}359}
363 360 
364template <typename Tx, typename Ty>361template <typename Tx, typename Ty>
365-__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::MainBlockCastSubExpVF(__local_mem__ float*& xFp32Ptr,362+__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::MainBlockCastSubExpVF(__ubuf__ float*& xFp32Ptr,
366- __local_mem__ Tx*& xPtr,363+ __ubuf__ Tx*& xPtr, __ubuf__ float*& xMaxPtr,
367- __local_mem__ float*& xMaxPtr, uint32_t a,364+ uint32_t a, uint32_t ubFactor)
368- uint32_t ubFactor)
369{365{
370 __VEC_SCOPE__366 __VEC_SCOPE__
371 {367 {
@@ -376,31 +372,30 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::MainBlockCastSubExpVF(__loc
376 uint32_t width = ubFactor;372 uint32_t width = ubFactor;
377 uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32);373 uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32);
378 374 
379- AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr);375+ AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr);
380 for (uint16_t j = 0; j < repeatTimes; j++) {376 for (uint16_t j = 0; j < repeatTimes; j++) {
381 mask = AscendC::MicroAPI::UpdateMask<float>(width);377 mask = AscendC::MicroAPI::UpdateMask<float>(width);
382 auto xAddr = xPtr + j * VL_FP32;378 auto xAddr = xPtr + j * VL_FP32;
383 auto xFp32Addr = xFp32Ptr + j * VL_FP32;379 auto xFp32Addr = xFp32Ptr + j * VL_FP32;
384 380 
385 if constexpr (xToFp32_) {381 if constexpr (xToFp32_) {
386- AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr);382+ AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr);
387 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask);383 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask);
388 } else {384 } else {
389- AscendC::MicroAPI::DataCopy(vreg1, xAddr);385+ AscendC::MicroAPI::LoadAlign(vreg1, xAddr);
390 }386 }
391 387 
392 AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask);388 AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask);
393 AscendC::MicroAPI::Exp(vreg3, vreg2, mask);389 AscendC::MicroAPI::Exp(vreg3, vreg2, mask);
394 390 
395- AscendC::MicroAPI::DataCopy(xFp32Addr, vreg3, mask);391+ AscendC::MicroAPI::StoreAlign(xFp32Addr, vreg3, mask);
396 }392 }
397 }393 }
398}394}
399 395 
400template <typename Tx, typename Ty>396template <typename Tx, typename Ty>
401-__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::FoldBlockCastSubExpVF(__local_mem__ float*& dstPtr,397+__aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::FoldBlockCastSubExpVF(__ubuf__ float*& dstPtr, __ubuf__ Tx*& xPtr,
402- __local_mem__ Tx*& xPtr,398+ __ubuf__ float*& xMaxPtr, uint32_t a,
403- __local_mem__ float*& xMaxPtr, uint32_t a,
404 uint32_t ubFactor)399 uint32_t ubFactor)
405{400{
406 __VEC_SCOPE__401 __VEC_SCOPE__
@@ -412,26 +407,26 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::FoldBlockCastSubExpVF(__loc
412 uint32_t width = ubFactor;407 uint32_t width = ubFactor;
413 uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32);408 uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32);
414 409 
415- AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr);410+ AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr);
416 for (uint16_t j = 0; j < repeatTimes; j++) {411 for (uint16_t j = 0; j < repeatTimes; j++) {
417 mask = AscendC::MicroAPI::UpdateMask<float>(width);412 mask = AscendC::MicroAPI::UpdateMask<float>(width);
418 auto xAddr = xPtr + j * VL_FP32;413 auto xAddr = xPtr + j * VL_FP32;
419 auto dstAddr = dstPtr + j * VL_FP32;414 auto dstAddr = dstPtr + j * VL_FP32;
420 415 
421 if constexpr (xToFp32_) {416 if constexpr (xToFp32_) {
422- AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr);417+ AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr);
423 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask);418 AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask);
424 } else {419 } else {
425- AscendC::MicroAPI::DataCopy(vreg1, xAddr);420+ AscendC::MicroAPI::LoadAlign(vreg1, xAddr);
426 }421 }
427 422 
428 AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask);423 AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask);
429 AscendC::MicroAPI::Exp(vreg3, vreg2, mask);424 AscendC::MicroAPI::Exp(vreg3, vreg2, mask);
430 425 
431- AscendC::MicroAPI::DataCopy(dstReg, dstAddr);426+ AscendC::MicroAPI::LoadAlign(dstReg, dstAddr);
432 AscendC::MicroAPI::Add(dstReg, dstReg, vreg3, mask);427 AscendC::MicroAPI::Add(dstReg, dstReg, vreg3, mask);
433 428 
434- AscendC::MicroAPI::DataCopy(dstAddr, dstReg, mask);429+ AscendC::MicroAPI::StoreAlign(dstAddr, dstReg, mask);
435 }430 }
436 }431 }
437}432}
@@ -454,9 +449,9 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::UpdateCache(const LocalTens
454 uint32_t outerLoopStride = VL_FP32;449 uint32_t outerLoopStride = VL_FP32;
455 uint32_t innerLoopStride = stride;450 uint32_t innerLoopStride = stride;
456 451 
457- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();452+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
458- __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride;453+ __ubuf__ float* cache = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheId * stride;
459- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();454+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
460 455 
461 __VEC_SCOPE__456 __VEC_SCOPE__
462 {457 {
@@ -465,16 +460,15 @@ __aicore__ inline void SoftmaxV2ArRecompute<Tx, Ty>::UpdateCache(const LocalTens
465 AscendC::MicroAPI::MaskReg pMask;460 AscendC::MicroAPI::MaskReg pMask;
466 for (uint16_t i = 0; i < outerLoopTimes; ++i) {461 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
467 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);462 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
468- AscendC::MicroAPI::DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride);463+ AscendC::MicroAPI::LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride);
469 for (uint16_t j = 0; j < innerLoopTimes; ++j) {464 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
470- AscendC::MicroAPI::DataCopy(bReg,465+ AscendC::MicroAPI::LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride);
471- (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride);
472 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);466 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
473 }467 }
474- AscendC::MicroAPI::DataCopy((__local_mem__ float*)cache + i * outerLoopStride, aReg, pMask);468+ AscendC::MicroAPI::StoreAlign((__ubuf__ float*)cache + i * outerLoopStride, aReg, pMask);
475 }469 }
476 }470 }
477}471}
478 472 
479} // namespace SoftmaxV2Ops473} // namespace SoftmaxV2Ops
480-#endif // SOFTMAX_V2_AR_RECOMPUTE_H474+#endif // SOFTMAX_V2_AR_RECOMPUTE_H
Mactivation/softmax_v2/op_kernel/arch35/softmax_v2_ar_small_r.h+20-20
@@ -105,9 +105,9 @@ private:
105 __aicore__ inline void CalcMaxSubExp(uint32_t curTileA0Len, uint32_t totalRLen)105 __aicore__ inline void CalcMaxSubExp(uint32_t curTileA0Len, uint32_t totalRLen)
106 {106 {
107 LocalTensor<Tx> xLocal_ = xQueue_.DeQue<Tx>();107 LocalTensor<Tx> xLocal_ = xQueue_.DeQue<Tx>();
108- __local_mem__ Tx* xAddr = (__local_mem__ Tx*)xLocal_.GetPhyAddr();108+ __ubuf__ Tx* xAddr = (__ubuf__ Tx*)xLocal_.GetPhyAddr();
109- __local_mem__ float* tmpAddr = (__local_mem__ float*)tmpLocal_.GetPhyAddr();109+ __ubuf__ float* tmpAddr = (__ubuf__ float*)tmpLocal_.GetPhyAddr();
110- __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr();110+ __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr();
111 111 
112 uint16_t aLoopTimes = ops::CeilDiv(curTileA0Len, VL_FP32);112 uint16_t aLoopTimes = ops::CeilDiv(curTileA0Len, VL_FP32);
113 uint16_t rLoopTimes = static_cast<uint16_t>(totalRLen);113 uint16_t rLoopTimes = static_cast<uint16_t>(totalRLen);
@@ -135,8 +135,8 @@ private:
135 LoadTensorForDtypeT(xAddr, reg2, mask, offset);135 LoadTensorForDtypeT(xAddr, reg2, mask, offset);
136 MicroAPI::Sub(reg2, reg2, maxReg, mask);136 MicroAPI::Sub(reg2, reg2, maxReg, mask);
137 MicroAPI::Exp(reg2, reg2, mask);137 MicroAPI::Exp(reg2, reg2, mask);
138- MicroAPI::DataCopy(tmpAddr + offset, reg2, mask);138+ MicroAPI::StoreAlign(tmpAddr + offset, reg2, mask);
139- MicroAPI::DataCopy(tmpAddr2 + offset, reg2, mask);139+ MicroAPI::StoreAlign(tmpAddr2 + offset, reg2, mask);
140 }140 }
141 }141 }
142 }142 }
@@ -152,10 +152,10 @@ private:
152 152 
153 __aicore__ inline void CalcOutput(uint32_t curTileA0Len, uint32_t totalRLen)153 __aicore__ inline void CalcOutput(uint32_t curTileA0Len, uint32_t totalRLen)
154 {154 {
155- __local_mem__ float* sumAddr = (__local_mem__ float*)sumLocal_.GetPhyAddr();155+ __ubuf__ float* sumAddr = (__ubuf__ float*)sumLocal_.GetPhyAddr();
156- __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr();156+ __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr();
157 tmpLocalTy_ = tmpLocal_.template ReinterpretCast<Ty>();157 tmpLocalTy_ = tmpLocal_.template ReinterpretCast<Ty>();
158- __local_mem__ Ty* tmpAddrTy = (__local_mem__ Ty*)tmpLocalTy_.GetPhyAddr();158+ __ubuf__ Ty* tmpAddrTy = (__ubuf__ Ty*)tmpLocalTy_.GetPhyAddr();
159 159 
160 uint16_t aLoopTimes = static_cast<uint16_t>(ops::CeilDiv(curTileA0Len, VL_FP32));160 uint16_t aLoopTimes = static_cast<uint16_t>(ops::CeilDiv(curTileA0Len, VL_FP32));
161 uint16_t rLoopTimes = static_cast<uint16_t>(tl_->totalRLen);161 uint16_t rLoopTimes = static_cast<uint16_t>(tl_->totalRLen);
@@ -171,21 +171,21 @@ private:
171 171 
172 for (uint16_t j = 0; j < aLoopTimes; j++) { // 列172 for (uint16_t j = 0; j < aLoopTimes; j++) { // 列
173 mask = MicroAPI::UpdateMask<float>(sreg);173 mask = MicroAPI::UpdateMask<float>(sreg);
174- MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(sumReg,174+ MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(sumReg,
175- (__local_mem__ float*)sumAddr + j * VL_FP32);175+ (__ubuf__ float*)sumAddr + j * VL_FP32);
176 176 
177 for (uint16_t i = 0; i < rLoopTimes; i++) { // 行177 for (uint16_t i = 0; i < rLoopTimes; i++) { // 行
178 uint32_t offset = j * VL_FP32 + i * tileA0LenLocal;178 uint32_t offset = j * VL_FP32 + i * tileA0LenLocal;
179 179 
180- MicroAPI::DataCopy(reg1, tmpAddr2 + offset);180+ MicroAPI::LoadAlign(reg1, tmpAddr2 + offset);
181 MicroAPI::Div(reg1, reg1, sumReg, mask);181 MicroAPI::Div(reg1, reg1, sumReg, mask);
182 182 
183 if constexpr (yToFp32_) {183 if constexpr (yToFp32_) {
184- MicroAPI::DataCopy(tmpAddrTy + offset, reg1, mask);184+ MicroAPI::StoreAlign(tmpAddrTy + offset, reg1, mask);
185 } else { // fp16、bf16185 } else { // fp16、bf16
186 MicroAPI::RegTensor<Ty> xFp16;186 MicroAPI::RegTensor<Ty> xFp16;
187 MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(xFp16, reg1, mask);187 MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(xFp16, reg1, mask);
188- MicroAPI::DataCopy<Ty, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + offset, xFp16, mask);188+ MicroAPI::StoreAlign<Ty, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + offset, xFp16, mask);
189 }189 }
190 }190 }
191 }191 }
@@ -257,29 +257,29 @@ private:
257 yQueue_.EnQue(yLocal);257 yQueue_.EnQue(yLocal);
258 }258 }
259 259 
260- __aicore__ inline void LoadTensorForDtypeT(const __local_mem__ Tx* src, RegTensor<float>& dst, MaskReg& preg,260+ __aicore__ inline void LoadTensorForDtypeT(const __ubuf__ Tx* src, RegTensor<float>& dst, MaskReg& preg,
261 uint32_t offset)261 uint32_t offset)
262 {262 {
263 if constexpr (xToFp32_) {263 if constexpr (xToFp32_) {
264 MicroAPI::RegTensor<Tx> xFp16;264 MicroAPI::RegTensor<Tx> xFp16;
265- MicroAPI::DataCopy<Tx, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ Tx*)src + offset));265+ MicroAPI::LoadAlign<Tx, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ Tx*)src + offset));
266 MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(dst, xFp16, preg);266 MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(dst, xFp16, preg);
267 } else {267 } else {
268- MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);268+ MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
269 }269 }
270 }270 }
271 271 
272 __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen)272 __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen)
273 {273 {
274- static constexpr MultiCopyConfig config = {false};274+ static constexpr NdDmaConfig config = {false};
275- MultiCopyLoopInfo<CONST_TWO> copyLoopInfo;275+ NdDmaLoopInfo<CONST_TWO> copyLoopInfo;
276 copyLoopInfo.loopSrcStride[0] = 1;276 copyLoopInfo.loopSrcStride[0] = 1;
277 copyLoopInfo.loopSrcStride[1] = totalRLen;277 copyLoopInfo.loopSrcStride[1] = totalRLen;
278 copyLoopInfo.loopDstStride[0] = tl_->tileA0Len;278 copyLoopInfo.loopDstStride[0] = tl_->tileA0Len;
279 copyLoopInfo.loopDstStride[1] = 1;279 copyLoopInfo.loopDstStride[1] = 1;
280 copyLoopInfo.loopSize[0] = totalRLen;280 copyLoopInfo.loopSize[0] = totalRLen;
281 copyLoopInfo.loopSize[1] = curTileA0Len;281 copyLoopInfo.loopSize[1] = curTileA0Len;
282- MultiCopyParams<Tx, CONST_TWO> params = {copyLoopInfo, 0};282+ NdDmaParams<Tx, CONST_TWO> params = {copyLoopInfo, 0};
283 283 
284 LocalTensor<Tx> xLocal_ = xQueue_.AllocTensor<Tx>();284 LocalTensor<Tx> xLocal_ = xQueue_.AllocTensor<Tx>();
285 DataCopy<Tx, CONST_TWO, config>(xLocal_, xGm_[xGmOffset], params);285 DataCopy<Tx, CONST_TWO, config>(xLocal_, xGm_[xGmOffset], params);
@@ -342,4 +342,4 @@ private:
342 342 
343} // namespace SoftmaxV2Ops343} // namespace SoftmaxV2Ops
344 344 
345-#endif345+#endif
Mactivation/softmax_v2/op_kernel/arch35/softmax_v2_ara_full_load.h+59-62
@@ -34,6 +34,8 @@ using AscendC::MicroAPI::MaskMergeMode;
34using AscendC::MicroAPI::MaskReg;34using AscendC::MicroAPI::MaskReg;
35using AscendC::MicroAPI::RegTensor;35using AscendC::MicroAPI::RegTensor;
36using AscendC::MicroAPI::StoreDist;36using AscendC::MicroAPI::StoreDist;
37+using AscendC::Reg::LoadAlign;
38+using AscendC::Reg::StoreAlign;
37 39 
38constexpr int64_t SCALE_COEF_TWO = 2;40constexpr int64_t SCALE_COEF_TWO = 2;
39constexpr int64_t SCALE_COEF_FOUR = 4;41constexpr int64_t SCALE_COEF_FOUR = 4;
@@ -135,10 +137,10 @@ private:
135 __aicore__ inline void Compute(int64_t curTileRLen, uint32_t curTileA0Len)137 __aicore__ inline void Compute(int64_t curTileRLen, uint32_t curTileA0Len)
136 {138 {
137 LocalTensor<T1> x = xQueue_.DeQue<T1>();139 LocalTensor<T1> x = xQueue_.DeQue<T1>();
138- __local_mem__ T1* xLocal = (__local_mem__ T1*)x.GetPhyAddr();140+ __ubuf__ T1* xLocal = (__ubuf__ T1*)x.GetPhyAddr();
139 141 
140 LocalTensor<float> xTmpTensor = xTmpBuf_.Get<float>();142 LocalTensor<float> xTmpTensor = xTmpBuf_.Get<float>();
141- __local_mem__ float* xTmpLocal = (__local_mem__ float*)xTmpTensor.GetPhyAddr();143+ __ubuf__ float* xTmpLocal = (__ubuf__ float*)xTmpTensor.GetPhyAddr();
142 144 
143 uint16_t loopA0Num = ops::CeilDiv(curTileA0Len, VL_FP32);145 uint16_t loopA0Num = ops::CeilDiv(curTileA0Len, VL_FP32);
144 146 
@@ -146,10 +148,10 @@ private:
146 xQueue_.FreeTensor<T1>(x);148 xQueue_.FreeTensor<T1>(x);
147 149 
148 LocalTensor<float> y = yQueue_.AllocTensor<float>();150 LocalTensor<float> y = yQueue_.AllocTensor<float>();
149- __local_mem__ float* yLocal = (__local_mem__ float*)y.GetPhyAddr();151+ __ubuf__ float* yLocal = (__ubuf__ float*)y.GetPhyAddr();
150 152 
151 LocalTensor<float> xReduceTensor = xReduceBuf_.Get<float>();153 LocalTensor<float> xReduceTensor = xReduceBuf_.Get<float>();
152- __local_mem__ float* xReduceLocal = (__local_mem__ float*)xReduceTensor.GetPhyAddr();154+ __ubuf__ float* xReduceLocal = (__ubuf__ float*)xReduceTensor.GetPhyAddr();
153 155 
154 VFReduceSum(xReduceLocal, xTmpLocal, yLocal, curTileRLen, curTileA0Len);156 VFReduceSum(xReduceLocal, xTmpLocal, yLocal, curTileRLen, curTileA0Len);
155 157 
@@ -158,7 +160,7 @@ private:
158 yQueue_.EnQue(y);160 yQueue_.EnQue(y);
159 }161 }
160 162 
161- __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,
162 uint16_t curTileA0Len, uint16_t loopA0Num)164 uint16_t curTileA0Len, uint16_t loopA0Num)
163 {165 {
164 uint32_t tileA0Len = tilingData_->tileA0Len;166 uint32_t tileA0Len = tilingData_->tileA0Len;
@@ -186,14 +188,14 @@ private:
186 LoadTensorForDtypeT1(xLocal, x, pregMask, xOffset1);188 LoadTensorForDtypeT1(xLocal, x, pregMask, xOffset1);
187 Sub(x, x, maxReg, pregMask);189 Sub(x, x, maxReg, pregMask);
188 Exp(x, x, pregMask);190 Exp(x, x, pregMask);
189- DataCopy(((__local_mem__ float*)xTmpLocal) + xOffset1, x, pregMask);191+ StoreAlign(((__ubuf__ float*)xTmpLocal) + xOffset1, x, pregMask);
190 }192 }
191 }193 }
192 }194 }
193 }195 }
194 196 
195- __aicore__ inline void VFReduceSum(__local_mem__ float* xReduceLocal, __local_mem__ float* xTmpLocal,197+ __aicore__ inline void VFReduceSum(__ubuf__ float* xReduceLocal, __ubuf__ float* xTmpLocal, __ubuf__ float* yInUb,
196- __local_mem__ float* yInUb, uint16_t curTileRLen, uint16_t curTileA0Len)198+ uint16_t curTileRLen, uint16_t curTileA0Len)
197 {199 {
198 if (tilingData_->totalRLen <= SCALE_COEF_TWO) {200 if (tilingData_->totalRLen <= SCALE_COEF_TWO) {
199 SumRLessThan2(xTmpLocal, xReduceLocal, curTileA0Len);201 SumRLessThan2(xTmpLocal, xReduceLocal, curTileA0Len);
@@ -206,8 +208,7 @@ private:
206 }208 }
207 }209 }
208 210 
209- __aicore__ inline void SumRLessThan2(__local_mem__ float* xTmpLocal, __local_mem__ float* xReduceLocal,211+ __aicore__ inline void SumRLessThan2(__ubuf__ float* xTmpLocal, __ubuf__ float* xReduceLocal, uint32_t curTileA0Len)
210- uint32_t curTileA0Len)
211 {212 {
212 uint32_t rStride = tilingData_->tileA0Len;213 uint32_t rStride = tilingData_->tileA0Len;
213 uint16_t rLoopCount = tilingData_->totalRLen;214 uint16_t rLoopCount = tilingData_->totalRLen;
@@ -224,16 +225,15 @@ private:
224 pregLoop = UpdateMask<float>(sreg0);225 pregLoop = UpdateMask<float>(sreg0);
225 Duplicate(sum, 0.0, pregLoop);226 Duplicate(sum, 0.0, pregLoop);
226 for (uint16_t i = 0; i < rLoopCount; i++) {227 for (uint16_t i = 0; i < rLoopCount; i++) {
227- DataCopy(xld, ((__local_mem__ float*)xTmpLocal + i * rStride + k * VL_FP32));228+ LoadAlign(xld, ((__ubuf__ float*)xTmpLocal + i * rStride + k * VL_FP32));
228 Add(sum, sum, xld, pregLoop);229 Add(sum, sum, xld, pregLoop);
229 }230 }
230- DataCopy(((__local_mem__ float*)xReduceLocal + k * VL_FP32), sum, pregLoop);231+ StoreAlign(((__ubuf__ float*)xReduceLocal + k * VL_FP32), sum, pregLoop);
231 }232 }
232 }233 }
233 }234 }
234 235 
235- __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)
236- uint32_t curTileA0Len)
237 {237 {
238 uint32_t remainderOffset = SCALE_COEF_TWO * tilingData_->tileA0Len;238 uint32_t remainderOffset = SCALE_COEF_TWO * tilingData_->tileA0Len;
239 uint32_t aLength = tilingData_->tileA0Len;239 uint32_t aLength = tilingData_->tileA0Len;
@@ -260,31 +260,30 @@ private:
260 for (uint16_t k = 0; k < aLoopCount; k++) {260 for (uint16_t k = 0; k < aLoopCount; k++) {
261 pregLoop = UpdateMask<float>(sreg0);261 pregLoop = UpdateMask<float>(sreg0);
262 uint32_t aLoopOffset = k * VL_FP32;262 uint32_t aLoopOffset = k * VL_FP32;
263- DataCopy(((__local_mem__ float*)xTmpLocal + validNumInXUb + aLoopOffset), zero, pregLoop);263+ StoreAlign(((__ubuf__ float*)xTmpLocal + validNumInXUb + aLoopOffset), zero, pregLoop);
264 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();264 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
265 TwoRowAddWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset,265 TwoRowAddWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset,
266 aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow);266 aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow);
267- DataCopy(((__local_mem__ float*)xReduceLocal + aLoopOffset), x1, pregLoop);267+ StoreAlign(((__ubuf__ float*)xReduceLocal + aLoopOffset), x1, pregLoop);
268 }268 }
269 }269 }
270 }270 }
271 271 
272- __aicore__ inline void TwoRowAddWithTail(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg,272+ __aicore__ inline void TwoRowAddWithTail(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg,
273 uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4,273 uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4,
274 RegTensor<float>& rem, RegTensor<float>& nextRow,274 RegTensor<float>& rem, RegTensor<float>& nextRow,
275 RegTensor<float>& remNextRow)275 RegTensor<float>& remNextRow)
276 {276 {
277- DataCopy(dst, ((__local_mem__ float*)(input) + (offset1)));277+ LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1)));
278- DataCopy(rem, ((__local_mem__ float*)(input) + (offset2)));278+ LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2)));
279 Add(dst, dst, rem, preg);279 Add(dst, dst, rem, preg);
280- DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3)));280+ LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3)));
281- DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4)));281+ LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4)));
282 Add(nextRow, nextRow, remNextRow, preg);282 Add(nextRow, nextRow, remNextRow, preg);
283 Add(dst, dst, nextRow, preg);283 Add(dst, dst, nextRow, preg);
284 }284 }
285 285 
286- __aicore__ inline void SumRLessThan8(__local_mem__ float* xTmpLocal, __local_mem__ float* xReduceLocal,286+ __aicore__ inline void SumRLessThan8(__ubuf__ float* xTmpLocal, __ubuf__ float* xReduceLocal, uint32_t curTileA0Len)
287- uint32_t curTileA0Len)
288 {287 {
289 uint32_t remainderOffset = SCALE_COEF_FOUR * tilingData_->tileA0Len;288 uint32_t remainderOffset = SCALE_COEF_FOUR * tilingData_->tileA0Len;
290 uint32_t aLength = tilingData_->tileA0Len;289 uint32_t aLength = tilingData_->tileA0Len;
@@ -316,7 +315,7 @@ private:
316 for (uint16_t k = 0; k < aLoopCount; k++) {315 for (uint16_t k = 0; k < aLoopCount; k++) {
317 pregLoop = UpdateMask<float>(sreg0);316 pregLoop = UpdateMask<float>(sreg0);
318 uint32_t aLoopOffset = k * VL_FP32;317 uint32_t aLoopOffset = k * VL_FP32;
319- DataCopy(((__local_mem__ float*)xTmpLocal + validNumInXUb + aLoopOffset), zero, pregLoop);318+ StoreAlign(((__ubuf__ float*)xTmpLocal + validNumInXUb + aLoopOffset), zero, pregLoop);
320 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();319 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
321 TwoRowAddWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset,320 TwoRowAddWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset,
322 aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow);321 aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow);
@@ -324,13 +323,13 @@ private:
324 remainderTailOffset2 + aLoopOffset, ROW_THREE_OFFSET * aLength + aLoopOffset,323 remainderTailOffset2 + aLoopOffset, ROW_THREE_OFFSET * aLength + aLoopOffset,
325 remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow);324 remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow);
326 Add(x1, x1, x2, pregLoop);325 Add(x1, x1, x2, pregLoop);
327- DataCopy(((__local_mem__ float*)xReduceLocal + aLoopOffset), x1, pregLoop);326+ StoreAlign(((__ubuf__ float*)xReduceLocal + aLoopOffset), x1, pregLoop);
328 }327 }
329 }328 }
330 }329 }
331 330 
332- __aicore__ inline void SumRMoreThan8(__local_mem__ float* xInUb, __local_mem__ float* yInUb,331+ __aicore__ inline void SumRMoreThan8(__ubuf__ float* xInUb, __ubuf__ float* yInUb, __ubuf__ float* xReduceLocal,
333- __local_mem__ float* xReduceLocal, uint32_t curTileA0Len)332+ uint32_t curTileA0Len)
334 {333 {
335 uint16_t remainderLoopCount = tilingData_->remainderLoopCount;334 uint16_t remainderLoopCount = tilingData_->remainderLoopCount;
336 uint16_t remainderLoopCountTmp = remainderLoopCount - 1;335 uint16_t remainderLoopCountTmp = remainderLoopCount - 1;
@@ -379,7 +378,7 @@ private:
379 for (uint16_t k = 0; k < aLoopCount; k++) {378 for (uint16_t k = 0; k < aLoopCount; k++) {
380 pregLoop = UpdateMask<float>(sreg0);379 pregLoop = UpdateMask<float>(sreg0);
381 uint32_t aLoopOffset = k * VL_FP32;380 uint32_t aLoopOffset = k * VL_FP32;
382- DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop);381+ StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop);
383 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();382 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
384 // 前半部分与后半部分中，都为8行的部分383 // 前半部分与后半部分中，都为8行的部分
385 for (uint16_t i = 0; i < remainderLoopCountTmp; i++) {384 for (uint16_t i = 0; i < remainderLoopCountTmp; i++) {
@@ -399,7 +398,7 @@ private:
399 remOffset + ROW_SEVEN_OFFSET * aLength, rem, nextRow, remNextRow);398 remOffset + ROW_SEVEN_OFFSET * aLength, rem, nextRow, remNextRow);
400 Add(x3, x3, x4, pregLoop);399 Add(x3, x3, x4, pregLoop);
401 Add(x1, x1, x3, pregLoop);400 Add(x1, x1, x3, pregLoop);
402- DataCopy(((__local_mem__ float*)yInUb + i * aLength + aLoopOffset), x1, pregLoop);401+ StoreAlign(((__ubuf__ float*)yInUb + i * aLength + aLoopOffset), x1, pregLoop);
403 }402 }
404 // 前半部分为8行，后半部分可能不足8行403 // 前半部分为8行，后半部分可能不足8行
405 {404 {
@@ -421,8 +420,8 @@ private:
421 remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow);420 remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow);
422 Add(x3, x3, x4, pregLoop);421 Add(x3, x3, x4, pregLoop);
423 Add(x1, x1, x3, pregLoop);422 Add(x1, x1, x3, pregLoop);
424- DataCopy(((__local_mem__ float*)yInUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1,423+ StoreAlign(((__ubuf__ float*)yInUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1,
425- pregLoop);424+ pregLoop);
426 }425 }
427 // 剩余的前半部分，一次for循环，处理8行426 // 剩余的前半部分，一次for循环，处理8行
428 for (uint16_t i = 0; i < quotientLoopCount; i++) {427 for (uint16_t i = 0; i < quotientLoopCount; i++) {
@@ -437,27 +436,27 @@ private:
437 baseOffset + ROW_SEVEN_OFFSET * aLength, nextRow);436 baseOffset + ROW_SEVEN_OFFSET * aLength, nextRow);
438 Add(x3, x3, x4, pregLoop);437 Add(x3, x3, x4, pregLoop);
439 Add(x1, x1, x3, pregLoop);438 Add(x1, x1, x3, pregLoop);
440- DataCopy(((__local_mem__ float*)yInUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1,439+ StoreAlign(((__ubuf__ float*)yInUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1,
441- pregLoop);440+ pregLoop);
442 }441 }
443 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();442 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
444- BinaryAddVF((__local_mem__ float*)yInUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop,443+ BinaryAddVF((__ubuf__ float*)yInUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop,
445 binaryAddLastLoop, pregLoop, x1, x2, x3, x4);444 binaryAddLastLoop, pregLoop, x1, x2, x3, x4);
446- DataCopy(x1, ((__local_mem__ float*)yInUb + aLoopOffset));445+ LoadAlign(x1, ((__ubuf__ float*)yInUb + aLoopOffset));
447- DataCopy(((__local_mem__ float*)xReduceLocal + aLoopOffset), x1, pregLoop);446+ StoreAlign(((__ubuf__ float*)xReduceLocal + aLoopOffset), x1, pregLoop);
448 }447 }
449 }448 }
450 }449 }
451 450 
452- __aicore__ inline void TwoRowAdd(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, uint32_t offset1,451+ __aicore__ inline void TwoRowAdd(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, uint32_t offset1,
453 uint32_t offset2, RegTensor<float>& nextRow)452 uint32_t offset2, RegTensor<float>& nextRow)
454 {453 {
455- DataCopy(dst, ((__local_mem__ float*)(input) + (offset1)));454+ LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1)));
456- DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset2)));455+ LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset2)));
457 Add(dst, dst, nextRow, preg);456 Add(dst, dst, nextRow, preg);
458 }457 }
459 458 
460- __aicore__ inline void BinaryAddVF(__local_mem__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint32_t offset,459+ __aicore__ inline void BinaryAddVF(__ubuf__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint32_t offset,
461 uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop,460 uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop,
462 MaskReg& pregLoop, RegTensor<float>& x1, RegTensor<float>& x2,461 MaskReg& pregLoop, RegTensor<float>& x1, RegTensor<float>& x2,
463 RegTensor<float>& x3, RegTensor<float>& x4)462 RegTensor<float>& x3, RegTensor<float>& x4)
@@ -466,32 +465,31 @@ private:
466 for (uint16_t i = 0; i < binaryAddKLoop; i++) {465 for (uint16_t i = 0; i < binaryAddKLoop; i++) {
467 curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET;466 curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET;
468 for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) {467 for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) {
469- DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset));468+ LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset));
470- DataCopy(x2,469+ LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset));
471- ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset));
472 Add(x1, x1, x2, pregLoop);470 Add(x1, x1, x2, pregLoop);
473- DataCopy(x3, ((__local_mem__ float*)binaryAddTmpAddr +471+ LoadAlign(x3, ((__ubuf__ float*)binaryAddTmpAddr +
474- (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset));472+ (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset));
475- DataCopy(x4, ((__local_mem__ float*)binaryAddTmpAddr +473+ LoadAlign(x4, ((__ubuf__ float*)binaryAddTmpAddr +
476- (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset));474+ (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset));
477 Add(x3, x3, x4, pregLoop);475 Add(x3, x3, x4, pregLoop);
478 Add(x1, x1, x3, pregLoop);476 Add(x1, x1, x3, pregLoop);
479- DataCopy(((__local_mem__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop);477+ StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop);
480 }478 }
481 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();479 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
482 }480 }
483 for (uint16_t i = 0; i < binaryAddLastLoop; i++) {481 for (uint16_t i = 0; i < binaryAddLastLoop; i++) {
484- DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + offset));482+ LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + offset));
485- DataCopy(x2, ((__local_mem__ float*)binaryAddTmpAddr + rLoopStride + offset));483+ LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + rLoopStride + offset));
486 Add(x1, x1, x2, pregLoop);484 Add(x1, x1, x2, pregLoop);
487- DataCopy(((__local_mem__ float*)binaryAddTmpAddr + offset), x1, pregLoop);485+ StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + offset), x1, pregLoop);
488 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();486 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
489 }487 }
490 }488 }
491 489 
492- __aicore__ inline void VFCalculateOutput(__local_mem__ float* yLocal, __local_mem__ float* xTmpLocal,490+ __aicore__ inline void VFCalculateOutput(__ubuf__ float* yLocal, __ubuf__ float* xTmpLocal,
493- __local_mem__ float* xReduceLocal, uint16_t curTileRLen,491+ __ubuf__ float* xReduceLocal, uint16_t curTileRLen, uint16_t curTileA0Len,
494- uint16_t curTileA0Len, uint16_t loopA0Num)492+ uint16_t loopA0Num)
495 {493 {
496 uint32_t tileA0Len = tilingData_->tileA0Len;494 uint32_t tileA0Len = tilingData_->tileA0Len;
497 __VEC_SCOPE__495 __VEC_SCOPE__
@@ -505,35 +503,34 @@ private:
505 503 
506 for (uint16_t k = 0; k < loopA0Num; k++) {504 for (uint16_t k = 0; k < loopA0Num; k++) {
507 pregMask = UpdateMask<float>(sreg);505 pregMask = UpdateMask<float>(sreg);
508- DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xReduceLocal + k * VL_FP32);506+ LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xReduceLocal + k * VL_FP32);
509 507 
510 for (uint16_t i = 0; i < curTileRLen; i++) {508 for (uint16_t i = 0; i < curTileRLen; i++) {
511 uint32_t xOffset = i * tileA0Len + k * VL_FP32;509 uint32_t xOffset = i * tileA0Len + k * VL_FP32;
512 510 
513- DataCopy<float, LoadDist::DIST_NORM>(xReg, (__local_mem__ float*)xTmpLocal + xOffset);511+ LoadAlign<float, LoadDist::DIST_NORM>(xReg, (__ubuf__ float*)xTmpLocal + xOffset);
514 Div(yReg, xReg, sumReg, pregMask);512 Div(yReg, xReg, sumReg, pregMask);
515 513 
516 // copy out514 // copy out
517 if constexpr (IsSameType<T2, float>::value) {515 if constexpr (IsSameType<T2, float>::value) {
518- DataCopy(((__local_mem__ float*)yLocal) + xOffset, yReg, pregMask);516+ StoreAlign(((__ubuf__ float*)yLocal) + xOffset, yReg, pregMask);
519 } else { // fp16、bf16517 } else { // fp16、bf16
520 RegTensor<T2> xFp16;518 RegTensor<T2> xFp16;
521 Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask);519 Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask);
522- DataCopy<T2, StoreDist::DIST_PACK_B32>(((__local_mem__ T2*)yLocal) + xOffset, xFp16, pregMask);520+ StoreAlign<T2, StoreDist::DIST_PACK_B32>(((__ubuf__ T2*)yLocal) + xOffset, xFp16, pregMask);
523 }521 }
524 }522 }
525 }523 }
526 }524 }
527 }525 }
528 526 
529- __aicore__ inline void LoadTensorForDtypeT1(__local_mem__ T1* src, RegTensor<float>& dst, MaskReg& preg,527+ __aicore__ inline void LoadTensorForDtypeT1(__ubuf__ T1* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)
530- uint32_t offset)
531 {528 {
532 if constexpr (IsSameType<T1, float>::value) {529 if constexpr (IsSameType<T1, float>::value) {
533- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);530+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
534 } else { // fp16、bf16531 } else { // fp16、bf16
535 RegTensor<T1> xFp16;532 RegTensor<T1> xFp16;
536- DataCopy<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T1*)src + offset));533+ LoadAlign<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T1*)src + offset));
537 Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg);534 Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg);
538 }535 }
539 }536 }
Mactivation/softmax_v2/op_kernel/arch35/softmax_v2_ara_recompute.h+29-28
@@ -30,6 +30,8 @@ using AscendC::MicroAPI::MaskMergeMode;
30using AscendC::MicroAPI::MaskReg;30using AscendC::MicroAPI::MaskReg;
31using AscendC::MicroAPI::RegTensor;31using AscendC::MicroAPI::RegTensor;
32using AscendC::MicroAPI::StoreDist;32using AscendC::MicroAPI::StoreDist;
33+using AscendC::Reg::LoadAlign;
34+using AscendC::Reg::StoreAlign;
33 35 
34template <typename T1, typename T2>36template <typename T1, typename T2>
35class SoftmaxV2ARARecompute : public SoftmaxV2OpsBase {37class SoftmaxV2ARARecompute : public SoftmaxV2OpsBase {
@@ -110,7 +112,7 @@ private:
110 {112 {
111 // max 初始化113 // max 初始化
112 LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>();114 LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>();
113- __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr();115+ __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr();
114 116 
115 __VEC_SCOPE__117 __VEC_SCOPE__
116 {118 {
@@ -120,7 +122,7 @@ private:
120 for (uint16_t k = 0; k < loopA0Num; k++) {122 for (uint16_t k = 0; k < loopA0Num; k++) {
121 pregMask = UpdateMask<float>(sreg);123 pregMask = UpdateMask<float>(sreg);
122 Duplicate(maxReg, static_cast<float>(-INFINITY), pregMask);124 Duplicate(maxReg, static_cast<float>(-INFINITY), pregMask);
123- DataCopy(((__local_mem__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask);125+ StoreAlign(((__ubuf__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask);
124 }126 }
125 }127 }
126 128 
@@ -134,7 +136,7 @@ private:
134 136 
135 CopyInX(xOffset, curTileRLen, curTileA0Len);137 CopyInX(xOffset, curTileRLen, curTileA0Len);
136 LocalTensor<T1> x = xQueue_.DeQue<T1>();138 LocalTensor<T1> x = xQueue_.DeQue<T1>();
137- __local_mem__ T1* xLocal = (__local_mem__ T1*)x.GetPhyAddr();139+ __ubuf__ T1* xLocal = (__ubuf__ T1*)x.GetPhyAddr();
138 140 
139 __VEC_SCOPE__141 __VEC_SCOPE__
140 {142 {
@@ -147,7 +149,7 @@ private:
147 for (uint16_t k = 0; k < loopA0Num; k++) {149 for (uint16_t k = 0; k < loopA0Num; k++) {
148 pregMask = UpdateMask<float>(sreg);150 pregMask = UpdateMask<float>(sreg);
149 // load max151 // load max
150- 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);
151 153 
152 for (uint16_t i = 0; i < curTileRLenVl; i++) {154 for (uint16_t i = 0; i < curTileRLenVl; i++) {
153 uint32_t offset = i * tileA0Len + k * VL_FP32;155 uint32_t offset = i * tileA0Len + k * VL_FP32;
@@ -155,7 +157,7 @@ private:
155 LoadTensorForDtypeT1(xLocal, x, pregMask, offset);157 LoadTensorForDtypeT1(xLocal, x, pregMask, offset);
156 Max(maxReg, maxReg, x, pregMask);158 Max(maxReg, maxReg, x, pregMask);
157 }159 }
158- DataCopy(((__local_mem__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask);160+ StoreAlign(((__ubuf__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask);
159 }161 }
160 }162 }
161 163 
@@ -212,9 +214,9 @@ private:
212 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);214 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);
213 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;215 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;
214 216 
215- __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr();217+ __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr();
216- __local_mem__ T1* src = (__local_mem__ T1*)xMain_.GetPhyAddr();218+ __ubuf__ T1* src = (__ubuf__ T1*)xMain_.GetPhyAddr();
217- __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr();219+ __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr();
218 220 
219 __VEC_SCOPE__221 __VEC_SCOPE__
220 {222 {
@@ -227,13 +229,13 @@ private:
227 229 
228 for (uint16_t j = 0; j < loopA0Num; ++j) {230 for (uint16_t j = 0; j < loopA0Num; ++j) {
229 pregMask = UpdateMask<float>(sreg);231 pregMask = UpdateMask<float>(sreg);
230- 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);
231 for (uint16_t i = 0; i < outerLoopTimes; ++i) {233 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
232 uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32;234 uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32;
233 LoadTensorForDtypeT1(src, srcReg, pregMask, xOffset);235 LoadTensorForDtypeT1(src, srcReg, pregMask, xOffset);
234 Sub(dstReg, srcReg, maxReg, pregMask);236 Sub(dstReg, srcReg, maxReg, pregMask);
235 Exp(dstReg, dstReg, pregMask);237 Exp(dstReg, dstReg, pregMask);
236- DataCopy((__local_mem__ float*)dst + xOffset, dstReg, pregMask);238+ StoreAlign((__ubuf__ float*)dst + xOffset, dstReg, pregMask);
237 }239 }
238 }240 }
239 }241 }
@@ -254,9 +256,9 @@ private:
254 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);256 uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen);
255 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;257 uint32_t outerLoopSrcStride = tilingData_->tileA0Len;
256 258 
257- __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr();259+ __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr();
258- __local_mem__ T1* src = (__local_mem__ T1*)xFold.GetPhyAddr();260+ __ubuf__ T1* src = (__ubuf__ T1*)xFold.GetPhyAddr();
259- __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr();261+ __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr();
260 262 
261 __VEC_SCOPE__263 __VEC_SCOPE__
262 {264 {
@@ -269,15 +271,15 @@ private:
269 271 
270 for (uint16_t j = 0; j < loopA0Num; ++j) {272 for (uint16_t j = 0; j < loopA0Num; ++j) {
271 pregMask = UpdateMask<float>(sreg);273 pregMask = UpdateMask<float>(sreg);
272- 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);
273 for (uint16_t i = 0; i < outerLoopTimes; ++i) {275 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
274 uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32;276 uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32;
275 LoadTensorForDtypeT1(src, srcReg, pregMask, i * outerLoopSrcStride + j * VL_FP32);277 LoadTensorForDtypeT1(src, srcReg, pregMask, i * outerLoopSrcStride + j * VL_FP32);
276 Sub(dstReg, srcReg, maxReg, pregMask);278 Sub(dstReg, srcReg, maxReg, pregMask);
277 Exp(dstReg, dstReg, pregMask);279 Exp(dstReg, dstReg, pregMask);
278- DataCopy(srcReg, (__local_mem__ float*)dst + xOffset);280+ LoadAlign(srcReg, (__ubuf__ float*)dst + xOffset);
279 Add(dstReg, dstReg, srcReg, pregMask);281 Add(dstReg, dstReg, srcReg, pregMask);
280- DataCopy((__local_mem__ float*)dst + xOffset, dstReg, pregMask);282+ StoreAlign((__ubuf__ float*)dst + xOffset, dstReg, pregMask);
281 }283 }
282 }284 }
283 }285 }
@@ -328,16 +330,16 @@ private:
328 __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)
329 {331 {
330 LocalTensor<T1> x = xQueue_.DeQue<T1>();332 LocalTensor<T1> x = xQueue_.DeQue<T1>();
331- __local_mem__ T1* xLocal = (__local_mem__ T1*)x.GetPhyAddr();333+ __ubuf__ T1* xLocal = (__ubuf__ T1*)x.GetPhyAddr();
332 334 
333 LocalTensor<T2> y = yQueue_.template AllocTensor<T2>();335 LocalTensor<T2> y = yQueue_.template AllocTensor<T2>();
334- __local_mem__ T2* yLocal = (__local_mem__ T2*)y.GetPhyAddr();336+ __ubuf__ T2* yLocal = (__ubuf__ T2*)y.GetPhyAddr();
335 337 
336 LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>();338 LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>();
337- __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr();339+ __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr();
338 340 
339 LocalTensor<float> xSumTensor = xSumBuf_.Get<float>();341 LocalTensor<float> xSumTensor = xSumBuf_.Get<float>();
340- __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor.GetPhyAddr();342+ __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor.GetPhyAddr();
341 343 
342 uint32_t tileA0Len = tilingData_->tileA0Len;344 uint32_t tileA0Len = tilingData_->tileA0Len;
343 uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen);345 uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen);
@@ -353,8 +355,8 @@ private:
353 355 
354 for (uint16_t k = 0; k < loopA0Num; k++) {356 for (uint16_t k = 0; k < loopA0Num; k++) {
355 pregMask = UpdateMask<float>(sreg);357 pregMask = UpdateMask<float>(sreg);
356- 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);
357- 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);
358 for (uint16_t i = 0; i < curTileRLenVl; i++) {360 for (uint16_t i = 0; i < curTileRLenVl; i++) {
359 uint32_t xOffset = i * tileA0Len + k * VL_FP32;361 uint32_t xOffset = i * tileA0Len + k * VL_FP32;
360 LoadTensorForDtypeT1(xLocal, xReg, pregMask, xOffset);362 LoadTensorForDtypeT1(xLocal, xReg, pregMask, xOffset);
@@ -365,11 +367,11 @@ private:
365 367 
366 // copy out368 // copy out
367 if constexpr (IsSameType<T2, float>::value) {369 if constexpr (IsSameType<T2, float>::value) {
368- DataCopy(((__local_mem__ float*)yLocal) + xOffset, yReg, pregMask);370+ StoreAlign(((__ubuf__ float*)yLocal) + xOffset, yReg, pregMask);
369 } else { // fp16、bf16371 } else { // fp16、bf16
370 RegTensor<T2> xFp16;372 RegTensor<T2> xFp16;
371 Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask);373 Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask);
372- 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);
373 }375 }
374 }376 }
375 }377 }
@@ -380,14 +382,13 @@ private:
380 xQueue_.FreeTensor<T1>(x);382 xQueue_.FreeTensor<T1>(x);
381 }383 }
382 384 
383- __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)
384- uint32_t offset)
385 {386 {
386 if constexpr (IsSameType<T1, float>::value) {387 if constexpr (IsSameType<T1, float>::value) {
387- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);388+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
388 } else { // fp16、bf16389 } else { // fp16、bf16
389 RegTensor<T1> xFp16;390 RegTensor<T1> xFp16;
390- DataCopy<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T1*)src + offset));391+ LoadAlign<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T1*)src + offset));
391 Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg);392 Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg);
392 }393 }
393 }394 }
Mactivation/softmax_v2/op_kernel/arch35/softmax_v2_base.h+170-167
@@ -24,6 +24,10 @@
24using namespace Ops::Base;24using namespace Ops::Base;
25namespace SoftmaxV2Ops {25namespace SoftmaxV2Ops {
26using namespace AscendC;26using namespace AscendC;
27+using AscendC::Reg::LoadAlign;
28+using AscendC::Reg::Move;
29+using AscendC::Reg::Reduce;
30+using AscendC::Reg::StoreAlign;
27 31 
28constexpr static AscendC::MicroAPI::CastTrait castTraitFp16ToFp32 = {32constexpr static AscendC::MicroAPI::CastTrait castTraitFp16ToFp32 = {
29 AscendC::MicroAPI::RegLayout::ZERO,33 AscendC::MicroAPI::RegLayout::ZERO,
@@ -176,8 +180,8 @@ __aicore__ inline void SoftmaxV2OpsBase::CastToFp32From(const LocalTensor<float>
176 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {180 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
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;
@@ -186,10 +190,10 @@ __aicore__ inline void SoftmaxV2OpsBase::CastToFp32From(const LocalTensor<float>
186 count = static_cast<uint32_t>(colSize);190 count = static_cast<uint32_t>(colSize);
187 for (uint16_t j = 0; j < innerLoopTimes; ++j) {191 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
188 pMask = AscendC::MicroAPI::UpdateMask<float>(count);192 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
189- DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(193+ LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
190- b16Reg, (__local_mem__ T*)src + i * outerLoopSrcStride + j * innerLoopStride);194+ b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + j * innerLoopStride);
191 Cast<float, T, castTraitFp16ToFp32>(fp32Reg, b16Reg, pMask);195 Cast<float, T, castTraitFp16ToFp32>(fp32Reg, b16Reg, pMask);
192- DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask);196+ StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask);
193 }197 }
194 }198 }
195 }199 }
@@ -219,8 +223,8 @@ __aicore__ inline void SoftmaxV2OpsBase::CastFromFp32To(const LocalTensor<T>& ds
219 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {223 if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
220 __VEC_SCOPE__224 __VEC_SCOPE__
221 {225 {
222- __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr();226+ __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr();
223- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();227+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
224 uint32_t count;228 uint32_t count;
225 AscendC::MicroAPI::RegTensor<float> fp32Reg;229 AscendC::MicroAPI::RegTensor<float> fp32Reg;
226 AscendC::MicroAPI::RegTensor<T> b16Reg;230 AscendC::MicroAPI::RegTensor<T> b16Reg;
@@ -229,10 +233,10 @@ __aicore__ inline void SoftmaxV2OpsBase::CastFromFp32To(const LocalTensor<T>& ds
229 count = static_cast<uint32_t>(colSize);233 count = static_cast<uint32_t>(colSize);
230 for (uint16_t j = 0; j < innerLoopTimes; ++j) {234 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
231 pMask = AscendC::MicroAPI::UpdateMask<float>(count);235 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
232- DataCopy(fp32Reg, (__local_mem__ float*)src + i * outerLoopSrcStride + j * innerLoopStride);236+ LoadAlign(fp32Reg, (__ubuf__ float*)src + i * outerLoopSrcStride + j * innerLoopStride);
233 Cast<T, float, castTraitFp32ToFp16>(b16Reg, fp32Reg, pMask);237 Cast<T, float, castTraitFp32ToFp16>(b16Reg, fp32Reg, pMask);
234- DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(238+ StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(
235- (__local_mem__ T*)dst + i * outerLoopDstStride + j * innerLoopStride, b16Reg, pMask);239+ (__ubuf__ T*)dst + i * outerLoopDstStride + j * innerLoopStride, b16Reg, pMask);
236 }240 }
237 }241 }
238 }242 }
@@ -316,19 +320,19 @@ __aicore__ inline void SoftmaxV2OpsBase::VectorAdd(const LocalTensor<float>& dst
316 uint16_t loopTimes = ops::CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize()));320 uint16_t loopTimes = ops::CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize()));
317 __VEC_SCOPE__321 __VEC_SCOPE__
318 {322 {
319- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();323+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
320- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();324+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
321- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();325+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
322 uint32_t sreg = static_cast<uint32_t>(count);326 uint32_t sreg = static_cast<uint32_t>(count);
323 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;327 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
324 AscendC::MicroAPI::MaskReg pMask;328 AscendC::MicroAPI::MaskReg pMask;
325 for (uint16_t i = 0; i < loopTimes; ++i) {329 for (uint16_t i = 0; i < loopTimes; ++i) {
326 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);330 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
327- DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32);331+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32);
328- DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32);332+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32);
329 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);333 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
330- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);334+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
331- DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask);335+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask);
332 }336 }
333 }337 }
334}338}
@@ -346,20 +350,20 @@ __aicore__ inline void SoftmaxV2OpsBase::VectorAdd(const LocalTensor<float>& dst
346 uint32_t innerLoopStride = stride;350 uint32_t innerLoopStride = stride;
347 __VEC_SCOPE__351 __VEC_SCOPE__
348 {352 {
349- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();353+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
350- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();354+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
351- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();355+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
352 uint32_t count = nSize;356 uint32_t count = nSize;
353 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;357 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
354 AscendC::MicroAPI::MaskReg pMask;358 AscendC::MicroAPI::MaskReg pMask;
355 for (uint16_t i = 0; i < outerLoopTimes; ++i) {359 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
356 pMask = AscendC::MicroAPI::UpdateMask<float>(count);360 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
357 for (uint16_t j = 0; j < innerLoopTimes; ++j) {361 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
358- DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride);362+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride);
359- DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + j * innerLoopStride);363+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + j * innerLoopStride);
360 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);364 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
361- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);365+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
362- DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask);366+ StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask);
363 }367 }
364 }368 }
365 }369 }
@@ -376,19 +380,19 @@ __aicore__ inline void SoftmaxV2OpsBase::VectorMul(const LocalTensor<float>& dst
376 uint16_t loopTimes = ops::CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize()));380 uint16_t loopTimes = ops::CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize()));
377 __VEC_SCOPE__381 __VEC_SCOPE__
378 {382 {
379- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();383+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
380- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();384+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
381- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();385+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
382 uint32_t sreg = static_cast<uint32_t>(count);386 uint32_t sreg = static_cast<uint32_t>(count);
383 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;387 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
384 AscendC::MicroAPI::MaskReg pMask;388 AscendC::MicroAPI::MaskReg pMask;
385 389 
386 for (uint16_t i = 0; i < loopTimes; ++i) {390 for (uint16_t i = 0; i < loopTimes; ++i) {
387 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);391 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
388- DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32);392+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32);
389- DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32);393+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32);
390 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);394 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
391- DataCopy((__local_mem__ float*)dst + i * VL_FP32, cReg, pMask);395+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, cReg, pMask);
392 }396 }
393 }397 }
394}398}
@@ -412,19 +416,19 @@ __aicore__ inline void SoftmaxV2OpsBase::NlastBroadcastMul(const LocalTensor<flo
412 uint32_t innerLoopStride = aSize;416 uint32_t innerLoopStride = aSize;
413 __VEC_SCOPE__417 __VEC_SCOPE__
414 {418 {
415- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();419+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
416- __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr();420+ __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr();
417- __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr();421+ __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr();
418 uint32_t count = static_cast<uint32_t>(aSize);422 uint32_t count = static_cast<uint32_t>(aSize);
419 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;423 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
420 AscendC::MicroAPI::MaskReg pMask;424 AscendC::MicroAPI::MaskReg pMask;
421 for (uint16_t i = 0; i < outerLoopTimes; ++i) {425 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
422 pMask = AscendC::MicroAPI::UpdateMask<float>(count);426 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
423- DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride);427+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride);
424 for (uint16_t j = 0; j < innerLoopTimes; ++j) {428 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
425- DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride);429+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride);
426 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);430 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
427- DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);431+ StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);
428 }432 }
429 }433 }
430 }434 }
@@ -449,40 +453,40 @@ __aicore__ inline void SoftmaxV2OpsBase::LastReduceSumSmallR(const LocalTensor<f
449 if (rSize <= VL_FP32) {453 if (rSize <= VL_FP32) {
450 __VEC_SCOPE__454 __VEC_SCOPE__
451 {455 {
452- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();456+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
453- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();457+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
454 uint32_t count = static_cast<uint32_t>(rSize);458 uint32_t count = static_cast<uint32_t>(rSize);
455 AscendC::MicroAPI::RegTensor<float> aReg, bReg;459 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
456 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);460 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
457- AscendC::MicroAPI::UnalignReg UReg;461+ AscendC::MicroAPI::UnalignRegForStore UReg;
458 for (uint16_t i = 0; i < loopTimes; ++i) {462 for (uint16_t i = 0; i < loopTimes; ++i) {
459- DataCopy(aReg, (__local_mem__ float*)src + i * stride);463+ LoadAlign(aReg, (__ubuf__ float*)src + i * stride);
460- ReduceSum(bReg, aReg, pMask);464+ Reduce<ReduceType::SUM>(bReg, aReg, pMask);
461- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);465+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
462 }466 }
463- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);467+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
464 }468 }
465 } else {469 } else {
466 __VEC_SCOPE__470 __VEC_SCOPE__
467 {471 {
468- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();472+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
469- __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr();473+ __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr();
470- __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32;474+ __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32;
471 uint32_t count = static_cast<uint32_t>(rSize - VL_FP32);475 uint32_t count = static_cast<uint32_t>(rSize - VL_FP32);
472 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;476 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
473- AscendC::MicroAPI::UnalignReg UReg;477+ AscendC::MicroAPI::UnalignRegForStore UReg;
474 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);478 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
475 AscendC::MicroAPI::MaskReg479 AscendC::MicroAPI::MaskReg
476 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();480 pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
477 for (uint16_t i = 0; i < loopTimes; ++i) {481 for (uint16_t i = 0; i < loopTimes; ++i) {
478- DataCopy(aReg, (__local_mem__ float*)src0 + i * stride);482+ LoadAlign(aReg, (__ubuf__ float*)src0 + i * stride);
479- DataCopy(bReg, (__local_mem__ float*)src1 + i * stride);483+ LoadAlign(bReg, (__ubuf__ float*)src1 + i * stride);
480 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);484 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
481- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);485+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
482- ReduceSum(bReg, aReg, pFull);486+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
483- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);487+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
484 }488 }
485- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);489+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
486 }490 }
487 }491 }
488}492}
@@ -527,43 +531,43 @@ __aicore__ inline void SoftmaxV2OpsBase::LastReduceSum(const LocalTensor<float>&
527 531 
528 __VEC_SCOPE__532 __VEC_SCOPE__
529 {533 {
530- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr();534+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr();
531- __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr();535+ __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr();
532- __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset;536+ __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset;
533- __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset;537+ __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset;
534- __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset;538+ __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset;
535- __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset;539+ __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset;
536 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();540 AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
537- AscendC::MicroAPI::UnalignReg UReg;541+ AscendC::MicroAPI::UnalignRegForStore UReg;
538 542 
539 for (uint16_t i = 0; i < outerLoopTimes; ++i) {543 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
540- dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;544+ dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride;
541 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {545 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {
542 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg;546 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg;
543- DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride);547+ LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride);
544- DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride);548+ LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride);
545 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull);549 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull);
546- ReduceSum(dReg, cReg, pFull);550+ Reduce<ReduceType::SUM>(dReg, cReg, pFull);
547- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1);551+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1);
548 }552 }
549 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {553 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {
550 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);554 uint32_t count = static_cast<uint32_t>(tailFoldElemCount);
551 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;555 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
552 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);556 AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count);
553- DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride);557+ LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride);
554- DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride);558+ LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride);
555 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);559 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask);
556- Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);560+ Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask);
557- ReduceSum(bReg, aReg, pFull);561+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
558- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);562+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
559 }563 }
560 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {564 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {
561 AscendC::MicroAPI::RegTensor<float> aReg, bReg;565 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
562- DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride);566+ LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride);
563- ReduceSum(bReg, aReg, pFull);567+ Reduce<ReduceType::SUM>(bReg, aReg, pFull);
564- AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1);568+ AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1);
565 }569 }
566- AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0);570+ AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0);
567 }571 }
568 }572 }
569 LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride);573 LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride);
@@ -571,13 +575,13 @@ __aicore__ inline void SoftmaxV2OpsBase::LastReduceSum(const LocalTensor<float>&
571 575 
572template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1>576template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1>
573struct NlastDichotomyAdd {577struct NlastDichotomyAdd {
574- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,578+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
575- __local_mem__ float*& srcA, __local_mem__ float*& srcB,579+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
576- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)580+ uint32_t stride)
577 {581 {
578 AscendC::MicroAPI::RegTensor<float> aReg, bReg;582 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
579- __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO;583+ __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO;
580- __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO;584+ __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO;
581 if constexpr (TailCount <= 0) {585 if constexpr (TailCount <= 0) {
582 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,586 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,
583 stride * CONST_TWO);587 stride * CONST_TWO);
@@ -585,13 +589,13 @@ struct NlastDichotomyAdd {
585 }589 }
586 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);590 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
587 }591 }
588- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,592+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
589- __local_mem__ float*& srcA, __local_mem__ float*& srcB,593+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
590- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset)594+ uint32_t stride, uint32_t offset)
591 {595 {
592 AscendC::MicroAPI::RegTensor<float> aReg, bReg;596 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
593- __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO;597+ __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO;
594- __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO;598+ __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO;
595 if constexpr (TailCount <= 0) {599 if constexpr (TailCount <= 0) {
596 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,600 NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask,
597 stride * CONST_TWO, offset);601 stride * CONST_TWO, offset);
@@ -609,49 +613,49 @@ struct NlastDichotomyAdd {
609 613 
610template <int32_t TailCount, int32_t Index, int32_t Depth>614template <int32_t TailCount, int32_t Index, int32_t Depth>
611struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> {615struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> {
612- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,616+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
613- __local_mem__ float*& srcA, __local_mem__ float*& srcB,617+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
614- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)618+ uint32_t stride)
615 {619 {
616 AscendC::MicroAPI::RegTensor<float> aReg, bReg;620 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
617- DataCopy(aReg, (__local_mem__ float*)srcA);621+ LoadAlign(aReg, (__ubuf__ float*)srcA);
618- DataCopy(bReg, (__local_mem__ float*)srcB);622+ LoadAlign(bReg, (__ubuf__ float*)srcB);
619 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);623 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
620 }624 }
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, uint32_t offset)627+ uint32_t stride, uint32_t offset)
624 {628 {
625 if constexpr (TailCount <= 0) {629 if constexpr (TailCount <= 0) {
626 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;630 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
627- DataCopy(aReg, (__local_mem__ float*)srcA);631+ LoadAlign(aReg, (__ubuf__ float*)srcA);
628- DataCopy(bReg, (__local_mem__ float*)srcA + offset);632+ LoadAlign(bReg, (__ubuf__ float*)srcA + offset);
629 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);633 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
630- DataCopy(bReg, (__local_mem__ float*)srcB);634+ LoadAlign(bReg, (__ubuf__ float*)srcB);
631- DataCopy(cReg, (__local_mem__ float*)srcB + offset);635+ LoadAlign(cReg, (__ubuf__ float*)srcB + offset);
632 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);636 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);
633 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);637 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
634 } else {638 } else {
635 if constexpr (Index + Depth < TailCount) {639 if constexpr (Index + Depth < TailCount) {
636 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;640 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
637- DataCopy(aReg, (__local_mem__ float*)srcA);641+ LoadAlign(aReg, (__ubuf__ float*)srcA);
638- DataCopy(bReg, (__local_mem__ float*)srcA + offset);642+ LoadAlign(bReg, (__ubuf__ float*)srcA + offset);
639 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);643 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
640- DataCopy(bReg, (__local_mem__ float*)srcB);644+ LoadAlign(bReg, (__ubuf__ float*)srcB);
641- DataCopy(cReg, (__local_mem__ float*)srcB + offset);645+ LoadAlign(cReg, (__ubuf__ float*)srcB + offset);
642 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);646 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask);
643 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);647 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
644 } else if constexpr (Index < TailCount) {648 } else if constexpr (Index < TailCount) {
645 AscendC::MicroAPI::RegTensor<float> aReg, bReg;649 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
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 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);654 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
651 } else {655 } else {
652 AscendC::MicroAPI::RegTensor<float> aReg, bReg;656 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
653- DataCopy(aReg, (__local_mem__ float*)srcA);657+ LoadAlign(aReg, (__ubuf__ float*)srcA);
654- DataCopy(bReg, (__local_mem__ float*)srcB);658+ LoadAlign(bReg, (__ubuf__ float*)srcB);
655 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);659 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
656 }660 }
657 }661 }
@@ -660,25 +664,25 @@ struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> {
660 664 
661template <>665template <>
662struct NlastDichotomyAdd<CONST_TWO> {666struct NlastDichotomyAdd<CONST_TWO> {
663- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,667+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
664- __local_mem__ float*& srcA, __local_mem__ float*& srcB,668+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
665- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)669+ uint32_t stride)
666 {670 {
667 AscendC::MicroAPI::RegTensor<float> aReg, bReg;671 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
668- DataCopy(aReg, (__local_mem__ float*)srcA);672+ LoadAlign(aReg, (__ubuf__ float*)srcA);
669- DataCopy(bReg, (__local_mem__ float*)srcB);673+ LoadAlign(bReg, (__ubuf__ float*)srcB);
670 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);674 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask);
671 }675 }
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, uint32_t offset)678+ uint32_t stride, uint32_t offset)
675 {679 {
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 }688 }
@@ -686,11 +690,11 @@ struct NlastDichotomyAdd<CONST_TWO> {
686 690 
687template <>691template <>
688struct NlastDichotomyAdd<1> {692struct NlastDichotomyAdd<1> {
689- __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc,693+ __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA,
690- __local_mem__ float*& srcA, __local_mem__ float*& srcB,694+ __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask,
691- AscendC::MicroAPI::MaskReg& pMask, uint32_t stride)695+ uint32_t stride)
692 {696 {
693- DataCopy(acc, (__local_mem__ float*)srcA);697+ LoadAlign(acc, (__ubuf__ float*)srcA);
694 }698 }
695};699};
696 700 
@@ -704,32 +708,32 @@ __aicore__ inline void SoftmaxV2OpsBase::NlastReduceSumSmallR(const LocalTensor<
704 if constexpr (RSize == 1) {708 if constexpr (RSize == 1) {
705 __VEC_SCOPE__709 __VEC_SCOPE__
706 {710 {
707- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();711+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
708- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();712+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
709 uint32_t count = static_cast<uint32_t>(aSize);713 uint32_t count = static_cast<uint32_t>(aSize);
710 AscendC::MicroAPI::RegTensor<float> aReg;714 AscendC::MicroAPI::RegTensor<float> aReg;
711 AscendC::MicroAPI::MaskReg pMask;715 AscendC::MicroAPI::MaskReg pMask;
712 for (uint16_t i = 0; i < loopTimes; ++i) {716 for (uint16_t i = 0; i < loopTimes; ++i) {
713 pMask = AscendC::MicroAPI::UpdateMask<float>(count);717 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
714- DataCopy(aReg, (__local_mem__ float*)src + i * VL_FP32);718+ LoadAlign(aReg, (__ubuf__ float*)src + i * VL_FP32);
715- DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask);719+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask);
716 }720 }
717 }721 }
718 } else {722 } else {
719 __VEC_SCOPE__723 __VEC_SCOPE__
720 {724 {
721- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();725+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
722- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr();726+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr();
723- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride;727+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride;
724 uint32_t count = static_cast<uint32_t>(aSize);728 uint32_t count = static_cast<uint32_t>(aSize);
725 AscendC::MicroAPI::RegTensor<float> aReg;729 AscendC::MicroAPI::RegTensor<float> aReg;
726 AscendC::MicroAPI::MaskReg pMask;730 AscendC::MicroAPI::MaskReg pMask;
727 for (uint16_t i = 0; i < loopTimes; ++i) {731 for (uint16_t i = 0; i < loopTimes; ++i) {
728- __local_mem__ float* curSrcA = srcA + i * VL_FP32;732+ __ubuf__ float* curSrcA = srcA + i * VL_FP32;
729- __local_mem__ float* curSrcB = srcB + i * VL_FP32;733+ __ubuf__ float* curSrcB = srcB + i * VL_FP32;
730 pMask = AscendC::MicroAPI::UpdateMask<float>(count);734 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
731 NlastDichotomyAdd<RSize>::LoadAndAccumulate(aReg, curSrcA, curSrcB, pMask, stride);735 NlastDichotomyAdd<RSize>::LoadAndAccumulate(aReg, curSrcA, curSrcB, pMask, stride);
732- DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask);736+ StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask);
733 }737 }
734 }738 }
735 }739 }
@@ -801,35 +805,34 @@ __aicore__ inline void SoftmaxV2OpsBase::NlastReduceSumLargeR(const LocalTensor<
801 for (uint16_t i = 0; i < outerLoopTimes; ++i) {805 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
802 pMask = plt_b32(count, POST_UPDATE);806 pMask = plt_b32(count, POST_UPDATE);
803 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {807 for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) {
804- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() +808+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopStride +
805- i * outerLoopStride + j * innerLoopStride;809+ j * innerLoopStride;
806- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride +810+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride + j * srcStride;
807- j * srcStride;811+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride + i * outerLoopStride +
808- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride +812+ j * srcStride;
809- i * outerLoopStride + j * srcStride;
810 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, foldOffset);813 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, foldOffset);
811- DataCopy((__local_mem__ float*)dst, aReg, pMask);814+ StoreAlign((__ubuf__ float*)dst, aReg, pMask);
812 }815 }
813 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {816 for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) {
814- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() +817+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopStride +
815- i * outerLoopStride + mainFoldLoopTimes * innerLoopStride;818+ mainFoldLoopTimes * innerLoopStride;
816- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride +819+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride +
817- mainFoldLoopTimes * srcStride;820+ mainFoldLoopTimes * srcStride;
818- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride +821+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride + i * outerLoopStride +
819- i * outerLoopStride + mainFoldLoopTimes * srcStride;822+ mainFoldLoopTimes * srcStride;
820 NlastDichotomyAdd<COMPRESSION, TailCount>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride,823 NlastDichotomyAdd<COMPRESSION, TailCount>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride,
821 foldOffset);824 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 < unFoldLoopTimes; ++j) {827 for (uint16_t j = 0; j < unFoldLoopTimes; ++j) {
825- __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + unFoldDstOffset +828+ __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + unFoldDstOffset +
826- i * outerLoopStride + j * innerLoopStride;829+ i * outerLoopStride + j * innerLoopStride;
827- __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset +830+ __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + i * outerLoopStride +
828- i * outerLoopStride + j * srcStride;831+ j * srcStride;
829- __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + stride +832+ __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + stride +
830- i * outerLoopStride + j * srcStride;833+ i * outerLoopStride + j * srcStride;
831 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride);834 NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride);
832- DataCopy((__local_mem__ float*)dst, aReg, pMask);835+ StoreAlign((__ubuf__ float*)dst, aReg, pMask);
833 }836 }
834 }837 }
835 }838 }
@@ -889,20 +892,20 @@ __aicore__ inline void SoftmaxV2OpsBase::UpdateCache(const LocalTensor<float>& d
889 uint32_t innerLoopStride = stride;892 uint32_t innerLoopStride = stride;
890 __VEC_SCOPE__893 __VEC_SCOPE__
891 {894 {
892- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();895+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
893- __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride;896+ __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride;
894- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();897+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
895 uint32_t sreg = static_cast<uint32_t>(count);898 uint32_t sreg = static_cast<uint32_t>(count);
896 AscendC::MicroAPI::RegTensor<float> aReg, bReg;899 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
897 AscendC::MicroAPI::MaskReg pMask;900 AscendC::MicroAPI::MaskReg pMask;
898 for (uint16_t i = 0; i < outerLoopTimes; ++i) {901 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
899 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);902 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
900- DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride);903+ LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride);
901 for (uint16_t j = 0; j < innerLoopTimes; ++j) {904 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
902- DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride);905+ LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride);
903 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);906 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
904 }907 }
905- DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask);908+ StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask);
906 }909 }
907 }910 }
908}911}
@@ -921,27 +924,27 @@ __aicore__ inline void SoftmaxV2OpsBase::Normalize(const LocalTensor<float>& dst
921 uint32_t innerLoopStride = VL_FP32;924 uint32_t innerLoopStride = VL_FP32;
922 __VEC_SCOPE__925 __VEC_SCOPE__
923 {926 {
924- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();927+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
925- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();928+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
926- __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr();929+ __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr();
927- __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr();930+ __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr();
928 uint32_t count;931 uint32_t count;
929 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;932 AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg;
930 AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg;933 AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg;
931 AscendC::MicroAPI::MaskReg pMask;934 AscendC::MicroAPI::MaskReg pMask;
932 for (uint16_t i = 0; i < outerLoopTimes; ++i) {935 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
933 count = static_cast<uint32_t>(colSize);936 count = static_cast<uint32_t>(colSize);
934- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i);937+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i);
935- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i);938+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i);
936 for (uint16_t j = 0; j < innerLoopTimes; ++j) {939 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
937 pMask = AscendC::MicroAPI::UpdateMask<float>(count);940 pMask = AscendC::MicroAPI::UpdateMask<float>(count);
938- DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + j * innerLoopStride);941+ LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + j * innerLoopStride);
939 Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask);942 Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask);
940 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask);943 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask);
941- DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);944+ StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask);
942 }945 }
943 }946 }
944 }947 }
945}948}
946} // namespace SoftmaxV2Ops949} // namespace SoftmaxV2Ops
947-#endif950+#endif
Mloss/nll_loss/op_kernel/arch35/nll_loss_mix.h+31-31
@@ -200,10 +200,10 @@ public:
200 __aicore__ inline void ReduceSumInCore(const LocalTensor<float>& reduceBuf, const LocalTensor<float>& midRes,200 __aicore__ inline void ReduceSumInCore(const LocalTensor<float>& reduceBuf, const LocalTensor<float>& midRes,
201 uint32_t mainReduceLength, uint32_t tailReduceLength)201 uint32_t mainReduceLength, uint32_t tailReduceLength)
202 {202 {
203- __local_mem__ float* mainAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr();203+ __ubuf__ float* mainAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr();
204- __local_mem__ float* tailAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr(mainReduceLength);204+ __ubuf__ float* tailAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr(mainReduceLength);
205- __local_mem__ float* remainAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr(tailReduceLength);205+ __ubuf__ float* remainAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr(tailReduceLength);
206- __local_mem__ float* midResAddr = (__ubuf__ float*)midRes.GetPhyAddr();206+ __ubuf__ float* midResAddr = (__ubuf__ float*)midRes.GetPhyAddr();
207 207 
208 uint32_t tailLoop = tailReduceLength / reduceOnceSize_;208 uint32_t tailLoop = tailReduceLength / reduceOnceSize_;
209 uint32_t mainLoop = (mainReduceLength - tailReduceLength) / reduceOnceSize_;209 uint32_t mainLoop = (mainReduceLength - tailReduceLength) / reduceOnceSize_;
@@ -219,47 +219,47 @@ public:
219 219 
220 for (uint16_t i = 0; i < static_cast<uint16_t>(tailLoop); ++i) {220 for (uint16_t i = 0; i < static_cast<uint16_t>(tailLoop); ++i) {
221 pregLoop = AscendC::MicroAPI::UpdateMask<float>(tailReduceLength);221 pregLoop = AscendC::MicroAPI::UpdateMask<float>(tailReduceLength);
222- DataCopy(main1, mainAddr + i * DOUBLE * vfFloatNum_);222+ LoadAlign(main1, mainAddr + i * DOUBLE * vfFloatNum_);
223- DataCopy(main2, mainAddr + (i * DOUBLE + 1) * vfFloatNum_);223+ LoadAlign(main2, mainAddr + (i * DOUBLE + 1) * vfFloatNum_);
224- DataCopy(tail1, tailAddr + i * DOUBLE * vfFloatNum_);224+ LoadAlign(tail1, tailAddr + i * DOUBLE * vfFloatNum_);
225- DataCopy(tail2, tailAddr + (i * DOUBLE + 1) * vfFloatNum_);225+ LoadAlign(tail2, tailAddr + (i * DOUBLE + 1) * vfFloatNum_);
226 226 
227 Add(main1, main1, tail1, pregLoop);227 Add(main1, main1, tail1, pregLoop);
228 Add(main2, main2, tail2, pregLoop);228 Add(main2, main2, tail2, pregLoop);
229 Add(main1, main1, main2, pregLoop);229 Add(main1, main1, main2, pregLoop);
230- ReduceSum(res, main1, pregLoop);230+ Reduce<ReduceType::SUM>(res, main1, pregLoop);
231- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(midResAddr + i, res, pregMain);231+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(midResAddr + i, res, pregMain);
232 }232 }
233 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,233 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,
234 AscendC::MicroAPI::MemType::VEC_LOAD>();234 AscendC::MicroAPI::MemType::VEC_LOAD>();
235 for (uint16_t i = 0; i < static_cast<uint16_t>(mainLoop); ++i) {235 for (uint16_t i = 0; i < static_cast<uint16_t>(mainLoop); ++i) {
236 uint32_t sreg0 = mainReduceLength - tailReduceLength;236 uint32_t sreg0 = mainReduceLength - tailReduceLength;
237 pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0);237 pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0);
238- DataCopy(main1, remainAddr + i * DOUBLE * vfFloatNum_);238+ LoadAlign(main1, remainAddr + i * DOUBLE * vfFloatNum_);
239- DataCopy(main2, remainAddr + (i * DOUBLE + 1) * vfFloatNum_);239+ LoadAlign(main2, remainAddr + (i * DOUBLE + 1) * vfFloatNum_);
240 Add(main1, main1, main2, pregLoop);240 Add(main1, main1, main2, pregLoop);
241- ReduceSum(res, main1, pregLoop);241+ Reduce<ReduceType::SUM>(res, main1, pregLoop);
242- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(midResAddr + tailLoop + i, res,242+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(midResAddr + tailLoop + i, res,
243- pregMain);243+ pregMain);
244 }244 }
245 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,245 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,
246 AscendC::MicroAPI::MemType::VEC_LOAD>();246 AscendC::MicroAPI::MemType::VEC_LOAD>();
247 for (uint16_t i = 0; i < static_cast<uint16_t>(restLoop); ++i) {247 for (uint16_t i = 0; i < static_cast<uint16_t>(restLoop); ++i) {
248 uint32_t sreg0 = tailLoop + mainLoop;248 uint32_t sreg0 = tailLoop + mainLoop;
249 pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0);249 pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0);
250- DataCopy(main1, midResAddr + i * DOUBLE * vfFloatNum_);250+ LoadAlign(main1, midResAddr + i * DOUBLE * vfFloatNum_);
251- DataCopy(main2, midResAddr + (i * DOUBLE + 1) * vfFloatNum_);251+ LoadAlign(main2, midResAddr + (i * DOUBLE + 1) * vfFloatNum_);
252 Add(main1, main1, main2, pregLoop);252 Add(main1, main1, main2, pregLoop);
253- ReduceSum(res, main1, pregLoop);253+ Reduce<ReduceType::SUM>(res, main1, pregLoop);
254- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(midResAddr + i, res, pregMain);254+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(midResAddr + i, res, pregMain);
255 }255 }
256 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,256 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,
257 AscendC::MicroAPI::MemType::VEC_LOAD>();257 AscendC::MicroAPI::MemType::VEC_LOAD>();
258 {258 {
259 pregLoop = AscendC::MicroAPI::UpdateMask<float>(lengthBeforeLastReduce);259 pregLoop = AscendC::MicroAPI::UpdateMask<float>(lengthBeforeLastReduce);
260- DataCopy(main1, midResAddr);260+ LoadAlign(main1, midResAddr);
261- ReduceSum(res, main1, pregLoop);261+ Reduce<ReduceType::SUM>(res, main1, pregLoop);
262- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(mainAddr, res, pregMain);262+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(mainAddr, res, pregMain);
263 }263 }
264 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,264 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,
265 AscendC::MicroAPI::MemType::VEC_LOAD>();265 AscendC::MicroAPI::MemType::VEC_LOAD>();
@@ -269,8 +269,8 @@ public:
269 __aicore__ inline void ReduceSum256(const LocalTensor<float>& reduceBuf,269 __aicore__ inline void ReduceSum256(const LocalTensor<float>& reduceBuf,
270 const LocalTensor<float>& nextLevelReduceBuf, uint32_t idx)270 const LocalTensor<float>& nextLevelReduceBuf, uint32_t idx)
271 {271 {
272- __local_mem__ float* mainAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr();272+ __ubuf__ float* mainAddr = (__ubuf__ float*)reduceBuf.GetPhyAddr();
273- __local_mem__ float* outAddr = (__ubuf__ float*)nextLevelReduceBuf.GetPhyAddr();273+ __ubuf__ float* outAddr = (__ubuf__ float*)nextLevelReduceBuf.GetPhyAddr();
274 __VEC_SCOPE__274 __VEC_SCOPE__
275 {275 {
276 AscendC::MicroAPI::RegTensor<float> main1, main2, main3, main4, res;276 AscendC::MicroAPI::RegTensor<float> main1, main2, main3, main4, res;
@@ -278,16 +278,16 @@ public:
278 AscendC::MicroAPI::MaskReg pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0);278 AscendC::MicroAPI::MaskReg pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0);
279 AscendC::MicroAPI::MaskReg279 AscendC::MicroAPI::MaskReg
280 pregMain = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();280 pregMain = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
281- DataCopy(main1, mainAddr);281+ LoadAlign(main1, mainAddr);
282- DataCopy(main2, mainAddr + vfFloatNum_);282+ LoadAlign(main2, mainAddr + vfFloatNum_);
283- DataCopy(main3, mainAddr + DOUBLE * vfFloatNum_);283+ LoadAlign(main3, mainAddr + DOUBLE * vfFloatNum_);
284- DataCopy(main4, mainAddr + THIRD * vfFloatNum_);284+ LoadAlign(main4, mainAddr + THIRD * vfFloatNum_);
285 285 
286 Add(main1, main1, main2, pregLoop);286 Add(main1, main1, main2, pregLoop);
287 Add(main3, main3, main4, pregLoop);287 Add(main3, main3, main4, pregLoop);
288 Add(main1, main1, main3, pregLoop);288 Add(main1, main1, main3, pregLoop);
289- ReduceSum(res, main1, pregLoop);289+ Reduce<ReduceType::SUM>(res, main1, pregLoop);
290- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(outAddr + idx, res, pregMain);290+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(outAddr + idx, res, pregMain);
291 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,291 AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE,
292 AscendC::MicroAPI::MemType::VEC_LOAD>();292 AscendC::MicroAPI::MemType::VEC_LOAD>();
293 }293 }
@@ -649,4 +649,4 @@ private:
649 int64_t productOfCHW_{0};649 int64_t productOfCHW_{0};
650 int64_t productOfHW_{0};650 int64_t productOfHW_{0};
651};651};
652-#endif652+#endif
Mloss/sparse_softmax_cross_entropy_with_logits/op_kernel/arch35/sparse_softmax_cross_entropy_with_logits_full_load.h+26-26
@@ -27,9 +27,9 @@ using namespace AscendC;
27using AscendC::MicroAPI::LoadDist;27using AscendC::MicroAPI::LoadDist;
28using AscendC::MicroAPI::MaskPattern;28using AscendC::MicroAPI::MaskPattern;
29using AscendC::MicroAPI::MaskReg;29using AscendC::MicroAPI::MaskReg;
30-using AscendC::MicroAPI::MaskUnPack;
31using AscendC::MicroAPI::RegTensor;30using AscendC::MicroAPI::RegTensor;
32using AscendC::MicroAPI::StoreDist;31using AscendC::MicroAPI::StoreDist;
32+using AscendC::MicroAPI::UnPack;
33using AscendC::MicroAPI::UpdateMask;33using AscendC::MicroAPI::UpdateMask;
34 34 
35template <typename T1, typename T2, uint64_t schId, uint64_t db>35template <typename T1, typename T2, uint64_t schId, uint64_t db>
@@ -361,12 +361,12 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsFullLoad<T1, T2, schId
361 361 
362 for (uint16_t i = 0; i < aTimes; i++) {362 for (uint16_t i = 0; i < aTimes; i++) {
363 AscendC::MicroAPI::Duplicate(featuresReg, minVal);363 AscendC::MicroAPI::Duplicate(featuresReg, minVal);
364- AscendC::MicroAPI::DataCopy(featuresReg1, featuresAddr + i * rAlign + repeatTimes * vfLen);364+ AscendC::MicroAPI::LoadAlign(featuresReg1, featuresAddr + i * rAlign + repeatTimes * vfLen);
365 AscendC::MicroAPI::Max(featuresReg1, featuresReg, featuresReg1, preg);365 AscendC::MicroAPI::Max(featuresReg1, featuresReg, featuresReg1, preg);
366- AscendC::MicroAPI::Copy<T1, AscendC::MicroAPI::MaskMergeMode::MERGING>(featuresReg, featuresReg1, preg);366+ AscendC::MicroAPI::Move<T1, AscendC::MicroAPI::MaskMergeMode::MERGING>(featuresReg, featuresReg1, preg);
367 for (uint16_t j = 0; j < repeatTimes; j++) {367 for (uint16_t j = 0; j < repeatTimes; j++) {
368 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rAlign, j, vfLen);368 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rAlign, j, vfLen);
369- AscendC::MicroAPI::DataCopy(featuresReg1, featuresAddr1, offset);369+ AscendC::MicroAPI::LoadAlign(featuresReg1, featuresAddr1, offset);
370 AscendC::MicroAPI::Max(featuresReg, featuresReg1, featuresReg, pregMain);370 AscendC::MicroAPI::Max(featuresReg, featuresReg1, featuresReg, pregMain);
371 }371 }
372 if constexpr (sizeof(T1) == 2) {372 if constexpr (sizeof(T1) == 2) {
@@ -379,11 +379,11 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsFullLoad<T1, T2, schId
379 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresRegLowest32, featuresRegLowest, pregReduce);379 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresRegLowest32, featuresRegLowest, pregReduce);
380 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresRegHighest32, featuresRegHighest, pregReduce);380 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresRegHighest32, featuresRegHighest, pregReduce);
381 AscendC::MicroAPI::Max(maxRegTemp, featuresRegLowest32, featuresRegHighest32, pregReduce);381 AscendC::MicroAPI::Max(maxRegTemp, featuresRegLowest32, featuresRegHighest32, pregReduce);
382- AscendC::MicroAPI::ReduceMax(maxReg, maxRegTemp, pregReduce);382+ AscendC::MicroAPI::Reduce<ReduceType::MAX>(maxReg, maxRegTemp, pregReduce);
383 } else {383 } else {
384- AscendC::MicroAPI::ReduceMax(maxReg, featuresReg, pregReduce);384+ AscendC::MicroAPI::Reduce<ReduceType::MAX>(maxReg, featuresReg, pregReduce);
385 }385 }
386- DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(maxAddr + i, maxReg, mergePreg);386+ StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(maxAddr + i, maxReg, mergePreg);
387 }387 }
388 }388 }
389}389}
@@ -420,35 +420,35 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsFullLoad<T1, T2, schId
420 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);420 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
421 421 
422 for (uint16_t i = 0; i < aTimes; i++) {422 for (uint16_t i = 0; i < aTimes; i++) {
423- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg32, maxAddr + i);423+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg32, maxAddr + i);
424 for (uint16_t j = 0; j < repeatTimes; j++) {424 for (uint16_t j = 0; j < repeatTimes; j++) {
425 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T1>(i, rAlign, j, vfLen);425 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T1>(i, rAlign, j, vfLen);
426 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);426 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);
427 if constexpr (sizeof(T1) == 2) {427 if constexpr (sizeof(T1) == 2) {
428- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(428+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
429 featuresReg, featuresAddr, offsetT);429 featuresReg, featuresAddr, offsetT);
430 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresReg32, featuresReg, pregMain);430 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresReg32, featuresReg, pregMain);
431 } else {431 } else {
432- AscendC::MicroAPI::DataCopy(featuresReg32, featuresAddr, offset);432+ AscendC::MicroAPI::LoadAlign(featuresReg32, featuresAddr, offset);
433 }433 }
434 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, pregMain);434 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, pregMain);
435 AscendC::MicroAPI::Exp(temp1Reg, subReg, pregMain);435 AscendC::MicroAPI::Exp(temp1Reg, subReg, pregMain);
436- AscendC::MicroAPI::DataCopy(temp1Addr, temp1Reg, offset, pregMain);436+ AscendC::MicroAPI::StoreAlign(temp1Addr, temp1Reg, offset, pregMain);
437- AscendC::MicroAPI::DataCopy(subAddr, subReg, offset, pregMain);437+ AscendC::MicroAPI::StoreAlign(subAddr, subReg, offset, pregMain);
438 }438 }
439 439 
440 for (uint16_t k = 0; k < tailLoop; k++) {440 for (uint16_t k = 0; k < tailLoop; k++) {
441 if constexpr (sizeof(T1) == 2) {441 if constexpr (sizeof(T1) == 2) {
442- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(442+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
443 featuresReg, featuresAddr + i * rAlign + repeatTimes * vfLen);443 featuresReg, featuresAddr + i * rAlign + repeatTimes * vfLen);
444 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresReg32, featuresReg, preg);444 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(featuresReg32, featuresReg, preg);
445 } else {445 } else {
446- AscendC::MicroAPI::DataCopy(featuresReg32, featuresAddr + i * rAlign + repeatTimes * vfLen);446+ AscendC::MicroAPI::LoadAlign(featuresReg32, featuresAddr + i * rAlign + repeatTimes * vfLen);
447 }447 }
448 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, preg);448 AscendC::MicroAPI::Sub(subReg, featuresReg32, maxReg32, preg);
449 AscendC::MicroAPI::Exp(temp1Reg, subReg, preg);449 AscendC::MicroAPI::Exp(temp1Reg, subReg, preg);
450- AscendC::MicroAPI::DataCopy(temp1Addr + i * rAlign + repeatTimes * vfLen, temp1Reg, pregAlign);450+ AscendC::MicroAPI::StoreAlign(temp1Addr + i * rAlign + repeatTimes * vfLen, temp1Reg, pregAlign);
451- AscendC::MicroAPI::DataCopy(subAddr + i * rAlign + repeatTimes * vfLen, subReg, preg);451+ AscendC::MicroAPI::StoreAlign(subAddr + i * rAlign + repeatTimes * vfLen, subReg, preg);
452 }452 }
453 }453 }
454 }454 }
@@ -494,31 +494,31 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsFullLoad<T1, T2, schId
494 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);494 AscendC::MicroAPI::MaskReg pregAlign = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
495 495 
496 for (uint16_t i = 0; i < aTimes; i++) {496 for (uint16_t i = 0; i < aTimes; i++) {
497- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(sumReg, sumAddr + i);497+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(sumReg, sumAddr + i);
498 for (uint16_t j = 0; j < repeatTimes; j++) {498 for (uint16_t j = 0; j < repeatTimes; j++) {
499 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T1>(i, rAlign, j, vfLen);499 AscendC::MicroAPI::AddrReg offsetT = AscendC::MicroAPI::CreateAddrReg<T1>(i, rAlign, j, vfLen);
500 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);500 AscendC::MicroAPI::AddrReg offset = AscendC::MicroAPI::CreateAddrReg<float>(i, rAlign, j, vfLen);
501- AscendC::MicroAPI::DataCopy(temp1Reg, temp1Addr, offset);501+ AscendC::MicroAPI::LoadAlign(temp1Reg, temp1Addr, offset);
502- AscendC::MicroAPI::DataCopy(subReg, subAddr, offset);502+ AscendC::MicroAPI::LoadAlign(subReg, subAddr, offset);
503 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, pregMain);503 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, pregMain);
504- AscendC::MicroAPI::DataCopy(temp1Addr, temp1Reg, offset, pregMain);504+ AscendC::MicroAPI::StoreAlign(temp1Addr, temp1Reg, offset, pregMain);
505 AscendC::MicroAPI::Log(logReg, sumReg, pregMain);505 AscendC::MicroAPI::Log(logReg, sumReg, pregMain);
506 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, pregMain);506 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, pregMain);
507- AscendC::MicroAPI::DataCopy(temp2Addr, temp2Reg, offset, pregMain);507+ AscendC::MicroAPI::StoreAlign(temp2Addr, temp2Reg, offset, pregMain);
508 }508 }
509 509 
510 for (uint16_t k = 0; k < tailLoop; k++) {510 for (uint16_t k = 0; k < tailLoop; k++) {
511- AscendC::MicroAPI::DataCopy(temp1Reg, temp1Addr + i * rAlign + repeatTimes * vfLen);511+ AscendC::MicroAPI::LoadAlign(temp1Reg, temp1Addr + i * rAlign + repeatTimes * vfLen);
512- AscendC::MicroAPI::DataCopy(subReg, subAddr + i * rAlign + repeatTimes * vfLen);512+ AscendC::MicroAPI::LoadAlign(subReg, subAddr + i * rAlign + repeatTimes * vfLen);
513 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, preg);513 AscendC::MicroAPI::Div(temp1Reg, temp1Reg, sumReg, preg);
514- AscendC::MicroAPI::DataCopy(temp1Addr + i * rAlign + repeatTimes * vfLen, temp1Reg, preg);514+ AscendC::MicroAPI::StoreAlign(temp1Addr + i * rAlign + repeatTimes * vfLen, temp1Reg, preg);
515 AscendC::MicroAPI::Log(logReg, sumReg, preg);515 AscendC::MicroAPI::Log(logReg, sumReg, preg);
516 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, preg);516 AscendC::MicroAPI::Sub(temp2Reg, logReg, subReg, preg);
517- AscendC::MicroAPI::DataCopy(temp2Addr + i * rAlign + repeatTimes * vfLen, temp2Reg, preg);517+ AscendC::MicroAPI::StoreAlign(temp2Addr + i * rAlign + repeatTimes * vfLen, temp2Reg, preg);
518 }518 }
519 }519 }
520 }520 }
521}521}
522 522 
523} // namespace SparseSoftmaxCrossEntropyWithLogits523} // namespace SparseSoftmaxCrossEntropyWithLogits
524-#endif524+#endif
Mloss/sparse_softmax_cross_entropy_with_logits/op_kernel/arch35/sparse_softmax_cross_entropy_with_logits_split_r.h+65-64
@@ -119,14 +119,14 @@ private:
119 int64_t tailBlockFactor; // a轴分核，尾核数据量119 int64_t tailBlockFactor; // a轴分核，尾核数据量
120 int64_t rUbNumFactor; // R轴切分，一次UB可以放下的数据量，全载模板下等于r，注意32b对齐120 int64_t rUbNumFactor; // R轴切分，一次UB可以放下的数据量，全载模板下等于r，注意32b对齐
121 int64_t aUbNumFactor; // A轴切分，一次UB可以放下的数据量，非全载模板下等于1，注意32b对齐121 int64_t aUbNumFactor; // A轴切分，一次UB可以放下的数据量，非全载模板下等于1，注意32b对齐
122- int64_t aLoopTimes; //主核A方向循环搬移数据的次数122+ int64_t aLoopTimes; // 主核A方向循环搬移数据的次数
123- int64_t aLoopTimesT; //尾核A方向循环搬移数据的次数123+ int64_t aLoopTimesT; // 尾核A方向循环搬移数据的次数
124- int64_t aLoopTail; //主核A方向尾块的数据量124+ int64_t aLoopTail; // 主核A方向尾块的数据量
125- int64_t aLoopTailT; //尾核A方向尾块的数据量125+ int64_t aLoopTailT; // 尾核A方向尾块的数据量
126- int64_t rLoopTime; //不能全载时，R轴反向的循环次数126+ int64_t rLoopTime; // 不能全载时，R轴反向的循环次数
127- int64_t rLoopTile; //不能全载时，R轴反向的尾块数据量127+ int64_t rLoopTile; // 不能全载时，R轴反向的尾块数据量
128- int64_t kTimesTail; //不能全载时，完全二分累加，存在主尾块相加的次数128+ int64_t kTimesTail; // 不能全载时，完全二分累加，存在主尾块相加的次数
129- int64_t kTimes; //不能全载时，完全二分累加，2的k次方内循环次数129+ int64_t kTimes; // 不能全载时，完全二分累加，2的k次方内循环次数
130 int64_t updateStart_;130 int64_t updateStart_;
131 int64_t rLoopTileAlign;131 int64_t rLoopTileAlign;
132 int64_t coreStartOffset;132 int64_t coreStartOffset;
@@ -299,17 +299,17 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
299 for (uint16_t j = 0; j < repeatTimes1; j++) {299 for (uint16_t j = 0; j < repeatTimes1; j++) {
300 preg = AscendC::MicroAPI::UpdateMask<float>(nTimes);300 preg = AscendC::MicroAPI::UpdateMask<float>(nTimes);
301 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLenfp32);301 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLenfp32);
302- AscendC::MicroAPI::DataCopy(srcReg1, maxUbAddr, srcOffset);302+ AscendC::MicroAPI::LoadAlign(srcReg1, maxUbAddr, srcOffset);
303 if constexpr (sizeof(T1) == 4) {303 if constexpr (sizeof(T1) == 4) {
304- AscendC::MicroAPI::DataCopy(srcReg2, maxUbOnceAddrB32, srcOffset);304+ AscendC::MicroAPI::LoadAlign(srcReg2, maxUbOnceAddrB32, srcOffset);
305 } else {305 } else {
306 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<T1>(j, vfLenfp32);306 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<T1>(j, vfLenfp32);
307- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg2B16, maxUbOnceAddr,307+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
308- srcOffset1);308+ srcReg2B16, maxUbOnceAddr, srcOffset1);
309 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcReg2, srcReg2B16, preg);309 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcReg2, srcReg2B16, preg);
310 }310 }
311 AscendC::MicroAPI::Max(srcReg3, srcReg1, srcReg2, preg);311 AscendC::MicroAPI::Max(srcReg3, srcReg1, srcReg2, preg);
312- AscendC::MicroAPI::DataCopy(maxUbAddr, srcReg3, srcOffset, preg);312+ AscendC::MicroAPI::StoreAlign(maxUbAddr, srcReg3, srcOffset, preg);
313 }313 }
314 }314 }
315}315}
@@ -352,7 +352,7 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
352 preg = AscendC::MicroAPI::UpdateMask<float>(nSize);352 preg = AscendC::MicroAPI::UpdateMask<float>(nSize);
353 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);353 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);
354 AscendC::MicroAPI::Duplicate(srcReg0, minValueFp32_);354 AscendC::MicroAPI::Duplicate(srcReg0, minValueFp32_);
355- AscendC::MicroAPI::DataCopy(maxUbAddr, srcReg0, srcOffset, preg);355+ AscendC::MicroAPI::StoreAlign(maxUbAddr, srcReg0, srcOffset, preg);
356 }356 }
357 }357 }
358}358}
@@ -381,14 +381,14 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
381 for (uint16_t i = 0; i < outerLoopTimes; ++i) {381 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
382 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);382 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
383 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat);383 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat);
384- AscendC::MicroAPI::DataCopy(aReg, srcUbAddr, srcOffset);384+ AscendC::MicroAPI::LoadAlign(aReg, srcUbAddr, srcOffset);
385 for (uint16_t j = 0; j < innerLoopTimes; ++j) {385 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
386 AscendC::MicroAPI::AddrReg srcOffsetJ = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat, j,386 AscendC::MicroAPI::AddrReg srcOffsetJ = AscendC::MicroAPI::CreateAddrReg<float>(i, elementOneRepeat, j,
387 stride);387 stride);
388- AscendC::MicroAPI::DataCopy(bReg, dstUbAddr, srcOffsetJ);388+ AscendC::MicroAPI::LoadAlign(bReg, dstUbAddr, srcOffsetJ);
389 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);389 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
390 }390 }
391- AscendC::MicroAPI::DataCopy(cahUbAddr, aReg, srcOffset, pMask);391+ AscendC::MicroAPI::StoreAlign(cahUbAddr, aReg, srcOffset, pMask);
392 }392 }
393 }393 }
394}394}
@@ -417,22 +417,22 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
417 AscendC::MicroAPI::RegTensor<float> subReg;417 AscendC::MicroAPI::RegTensor<float> subReg;
418 AscendC::MicroAPI::RegTensor<float> expReg;418 AscendC::MicroAPI::RegTensor<float> expReg;
419 for (uint16_t i = 0; i < nTimes; i++) {419 for (uint16_t i = 0; i < nTimes; i++) {
420- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);420+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);
421 for (uint16_t j = 0; j < repeatTimes1; j++) {421 for (uint16_t j = 0; j < repeatTimes1; j++) {
422 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j, vfLen);422 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j, vfLen);
423 if constexpr (sizeof(T1) == 2) {423 if constexpr (sizeof(T1) == 2) {
424 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNum, j, vfLen);424 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNum, j, vfLen);
425- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,425+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
426- srcOffset);426+ srcOffset);
427 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);427 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
428 } else {428 } else {
429 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j,429 AscendC::MicroAPI::AddrReg srcOffset1 = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNum, j,
430 vfLen);430 vfLen);
431- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset1);431+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset1);
432 }432 }
433 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);433 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
434 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);434 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
435- AscendC::MicroAPI::DataCopy(outUbAddr, expReg, outOffset, copyOutReg);435+ AscendC::MicroAPI::StoreAlign(outUbAddr, expReg, outOffset, copyOutReg);
436 }436 }
437 }437 }
438 }438 }
@@ -469,38 +469,38 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
469 AscendC::MicroAPI::RegTensor<float> outReg1;469 AscendC::MicroAPI::RegTensor<float> outReg1;
470 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);470 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
471 for (uint16_t i = 0; i < nTimes; i++) {471 for (uint16_t i = 0; i < nTimes; i++) {
472- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);472+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);
473 for (uint16_t j = 0; j < repeatTimes1; j++) {473 for (uint16_t j = 0; j < repeatTimes1; j++) {
474 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumTAlign, j,474 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumTAlign, j,
475 vfLen);475 vfLen);
476 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,476 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
477 vfLen);477 vfLen);
478 if constexpr (sizeof(T1) == 2) {478 if constexpr (sizeof(T1) == 2) {
479- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,479+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
480- srcOffset);480+ srcOffset);
481 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);481 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
482 } else {482 } else {
483- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);483+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
484 }484 }
485 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);485 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
486 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);486 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
487- AscendC::MicroAPI::DataCopy(outReg, outUbAddr, outOffset);487+ AscendC::MicroAPI::LoadAlign(outReg, outUbAddr, outOffset);
488 AscendC::MicroAPI::Add(outReg1, expReg, outReg, copyOutReg);488 AscendC::MicroAPI::Add(outReg1, expReg, outReg, copyOutReg);
489- AscendC::MicroAPI::DataCopy(outUbAddr, outReg1, outOffset, copyOutReg);489+ AscendC::MicroAPI::StoreAlign(outUbAddr, outReg1, outOffset, copyOutReg);
490 }490 }
491 for (uint16_t k = 0; k < tailLoopTimes; k++) {491 for (uint16_t k = 0; k < tailLoopTimes; k++) {
492 if constexpr (sizeof(T1) == 2) {492 if constexpr (sizeof(T1) == 2) {
493- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(493+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
494 srcReg0, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);494 srcReg0, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
495 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, preg);495 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, preg);
496 } else {496 } else {
497- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);497+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumTAlign + repeatTimes1 * vfLen);
498 }498 }
499 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);499 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
500 AscendC::MicroAPI::Exp(expReg, subReg, preg);500 AscendC::MicroAPI::Exp(expReg, subReg, preg);
501- AscendC::MicroAPI::DataCopy(outReg, outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);501+ AscendC::MicroAPI::LoadAlign(outReg, outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
502 AscendC::MicroAPI::Add(outReg1, expReg, outReg, preg);502 AscendC::MicroAPI::Add(outReg1, expReg, outReg, preg);
503- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, outReg1, preg);503+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, outReg1, preg);
504 }504 }
505 }505 }
506 }506 }
@@ -536,33 +536,33 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
536 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);536 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
537 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);537 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
538 for (uint16_t i = 0; i < nTimes; i++) {538 for (uint16_t i = 0; i < nTimes; i++) {
539- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);539+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);
540 for (uint16_t j = 0; j < repeatTimes1; j++) {540 for (uint16_t j = 0; j < repeatTimes1; j++) {
541 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumAlign, j, vfLen);541 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumAlign, j, vfLen);
542 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,542 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
543 vfLen);543 vfLen);
544 if constexpr (sizeof(T1) == 2) {544 if constexpr (sizeof(T1) == 2) {
545- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,545+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
546- srcOffset);546+ srcOffset);
547 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);547 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
548 } else {548 } else {
549- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);549+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
550 }550 }
551 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);551 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
552 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);552 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
553- AscendC::MicroAPI::DataCopy(outUbAddr, expReg, outOffset, copyOutReg);553+ AscendC::MicroAPI::StoreAlign(outUbAddr, expReg, outOffset, copyOutReg);
554 }554 }
555 for (uint16_t k = 0; k < tailLoopTimes; k++) {555 for (uint16_t k = 0; k < tailLoopTimes; k++) {
556 if constexpr (sizeof(T1) == 2) {556 if constexpr (sizeof(T1) == 2) {
557- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(557+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
558 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);558 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
559 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, preg);559 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, preg);
560 } else {560 } else {
561- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);561+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
562 }562 }
563 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);563 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
564 AscendC::MicroAPI::Exp(expReg, subReg, preg);564 AscendC::MicroAPI::Exp(expReg, subReg, preg);
565- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg, preg1);565+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg, preg1);
566 }566 }
567 }567 }
568 }568 }
@@ -639,11 +639,11 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
639 for (uint16_t j = 0; j < repeatTimes1; j++) {639 for (uint16_t j = 0; j < repeatTimes1; j++) {
640 preg = AscendC::MicroAPI::UpdateMask<float>(nSize1);640 preg = AscendC::MicroAPI::UpdateMask<float>(nSize1);
641 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);641 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);
642- AscendC::MicroAPI::DataCopy(srcReg0, cacheUbAddr, srcOffset);642+ AscendC::MicroAPI::LoadAlign(srcReg0, cacheUbAddr, srcOffset);
643 AscendC::MicroAPI::Log(logReg0, srcReg0, preg);643 AscendC::MicroAPI::Log(logReg0, srcReg0, preg);
644- AscendC::MicroAPI::Copy(sumReg0, srcReg0, preg);644+ AscendC::MicroAPI::Move(sumReg0, srcReg0, preg);
645- AscendC::MicroAPI::DataCopy(logUbAddr, logReg0, srcOffset, preg);645+ AscendC::MicroAPI::StoreAlign(logUbAddr, logReg0, srcOffset, preg);
646- AscendC::MicroAPI::DataCopy(sumUbAddr, sumReg0, srcOffset, preg);646+ AscendC::MicroAPI::StoreAlign(sumUbAddr, sumReg0, srcOffset, preg);
647 }647 }
648 }648 }
649}649}
@@ -685,13 +685,13 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
685 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);685 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<float>(j, vfLen);
686 AscendC::MicroAPI::MaskReg686 AscendC::MicroAPI::MaskReg
687 regAllFp32 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();687 regAllFp32 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
688- AscendC::MicroAPI::DataCopy(srcReg0, cacheUbAddr, srcOffset);688+ AscendC::MicroAPI::LoadAlign(srcReg0, cacheUbAddr, srcOffset);
689 if constexpr (sizeof(T1) == 2) {689 if constexpr (sizeof(T1) == 2) {
690 AscendC::MicroAPI::Cast<T1, float, castB32ToB16>(lossReg0, srcReg0, preg);690 AscendC::MicroAPI::Cast<T1, float, castB32ToB16>(lossReg0, srcReg0, preg);
691- AscendC::MicroAPI::DataCopy<T1, StoreDist::DIST_PACK_B32>(lossUbAddr, lossReg0, preg);691+ AscendC::MicroAPI::StoreAlign<T1, StoreDist::DIST_PACK_B32>(lossUbAddr, lossReg0, preg);
692 } else {692 } else {
693- AscendC::MicroAPI::Copy(lossReg0, srcReg0, preg);693+ AscendC::MicroAPI::Move(lossReg0, srcReg0, preg);
694- AscendC::MicroAPI::DataCopy(lossUbAddr, lossReg0, srcOffset, preg);694+ AscendC::MicroAPI::StoreAlign(lossUbAddr, lossReg0, srcOffset, preg);
695 }695 }
696 }696 }
697 }697 }
@@ -774,54 +774,55 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
774 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);774 AscendC::MicroAPI::MaskReg preg = AscendC::MicroAPI::UpdateMask<float>(tailNum);
775 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);775 AscendC::MicroAPI::MaskReg preg1 = AscendC::MicroAPI::UpdateMask<float>(tailNumAlign);
776 for (uint16_t i = 0; i < nTimes; i++) {776 for (uint16_t i = 0; i < nTimes; i++) {
777- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);777+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(maxReg, maxUbAddr + i);
778- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(sumReg, sumUbAddr + i);778+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(sumReg, sumUbAddr + i);
779- AscendC::MicroAPI::DataCopy<float, LoadDist::DIST_BRC_B32>(logReg, logUbAddr + i);779+ AscendC::MicroAPI::LoadAlign<float, LoadDist::DIST_BRC_B32>(logReg, logUbAddr + i);
780 for (uint16_t j = 0; j < repeatTimes1; j++) {780 for (uint16_t j = 0; j < repeatTimes1; j++) {
781 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumAlign, j, vfLen);781 AscendC::MicroAPI::AddrReg srcOffset = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumAlign, j, vfLen);
782 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,782 AscendC::MicroAPI::AddrReg outOffset = AscendC::MicroAPI::CreateAddrReg<float>(i, rOnceNumAlign, j,
783 vfLen);783 vfLen);
784 if constexpr (sizeof(T1) == 2) {784 if constexpr (sizeof(T1) == 2) {
785- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,785+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(srcReg0, inputUbAddr,
786- srcOffset);786+ srcOffset);
787 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);787 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, copyOutReg);
788 } else {788 } else {
789- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr, srcOffset);789+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr, srcOffset);
790 }790 }
791 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);791 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, copyOutReg);
792 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);792 AscendC::MicroAPI::Exp(expReg, subReg, copyOutReg);
793 AscendC::MicroAPI::Div(expReg, expReg, sumReg, copyOutReg);793 AscendC::MicroAPI::Div(expReg, expReg, sumReg, copyOutReg);
794 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);794 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, copyOutReg);
795- AscendC::MicroAPI::DataCopy(outUbAddr, tmpReg, outOffset, copyOutReg);795+ AscendC::MicroAPI::StoreAlign(outUbAddr, tmpReg, outOffset, copyOutReg);
796 if constexpr (sizeof(T1) == 2) {796 if constexpr (sizeof(T1) == 2) {
797 AscendC::MicroAPI::AddrReg outOffset1 = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumAlign, j,797 AscendC::MicroAPI::AddrReg outOffset1 = AscendC::MicroAPI::CreateAddrReg<T1>(i, rOnceNumAlign, j,
798 vfLen);798 vfLen);
799 AscendC::MicroAPI::Cast<T1, float, castB32ToB16>(backProbReg, expReg, copyOutReg);799 AscendC::MicroAPI::Cast<T1, float, castB32ToB16>(backProbReg, expReg, copyOutReg);
800- AscendC::MicroAPI::DataCopy<T1, StoreDist::DIST_PACK_B32>(backProbAddr, backProbReg, outOffset1,800+ AscendC::MicroAPI::StoreAlign<T1, StoreDist::DIST_PACK_B32>(backProbAddr, backProbReg, outOffset1,
801- copyOutReg);801+ copyOutReg);
802 } else {802 } else {
803- AscendC::MicroAPI::DataCopy(backProbAddr, expReg, outOffset, copyOutReg);803+ AscendC::MicroAPI::StoreAlign(backProbAddr, expReg, outOffset, copyOutReg);
804 }804 }
805 }805 }
806 for (uint16_t k = 0; k < tailLoopTimes; k++) {806 for (uint16_t k = 0; k < tailLoopTimes; k++) {
807 if constexpr (sizeof(T1) == 2) {807 if constexpr (sizeof(T1) == 2) {
808- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(808+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(
809 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);809 srcReg0, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
810 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, preg);810 AscendC::MicroAPI::Cast<float, T1, castB16ToB32>(srcRegfp32, srcReg0, preg);
811 } else {811 } else {
812- AscendC::MicroAPI::DataCopy(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);812+ AscendC::MicroAPI::LoadAlign(srcRegfp32, inputUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen);
813 }813 }
814 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);814 AscendC::MicroAPI::Sub(subReg, srcRegfp32, maxReg, preg);
815 AscendC::MicroAPI::Exp(expReg, subReg, preg);815 AscendC::MicroAPI::Exp(expReg, subReg, preg);
816 AscendC::MicroAPI::Div(expReg, expReg, sumReg, preg);816 AscendC::MicroAPI::Div(expReg, expReg, sumReg, preg);
817 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, preg);817 AscendC::MicroAPI::Sub(tmpReg, logReg, subReg, preg);
818- AscendC::MicroAPI::DataCopy(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, tmpReg, preg);818+ AscendC::MicroAPI::StoreAlign(outUbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, tmpReg, preg);
819 if constexpr (sizeof(T1) == 2) {819 if constexpr (sizeof(T1) == 2) {
820 AscendC::MicroAPI::Cast<T1, float, castB32ToB16>(backProbReg, expReg, preg);820 AscendC::MicroAPI::Cast<T1, float, castB32ToB16>(backProbReg, expReg, preg);
821- AscendC::MicroAPI::DataCopy<T1, StoreDist::DIST_PACK_B32>(821+ AscendC::MicroAPI::StoreAlign<T1, StoreDist::DIST_PACK_B32>(
822 backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, backProbReg, preg);822 backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, backProbReg, preg);
823 } else {823 } else {
824- AscendC::MicroAPI::DataCopy(backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg, preg);824+ AscendC::MicroAPI::StoreAlign(backProbAddr + i * rOnceNumAlign + repeatTimes1 * vfLen, expReg,
825+ preg);
825 }826 }
826 }827 }
827 }828 }
@@ -848,4 +849,4 @@ __aicore__ inline void SparseSoftmaxCrossEntropyWithLogitsSplitR<T1, T2, schId,
848}849}
849 850 
850} // namespace SparseSoftmaxCrossEntropyWithLogits851} // namespace SparseSoftmaxCrossEntropyWithLogits
851-#endif // SPARSE_SOFTMAX_CROSS_ENTROPY_WITH_LOGITS_SPLIT_R_H852+#endif // SPARSE_SOFTMAX_CROSS_ENTROPY_WITH_LOGITS_SPLIT_R_H
Mnorm/rms_norm/op_kernel/arch35/rms_norm_regbase.h+1-1
@@ -185,7 +185,7 @@ private:
185 TQue<QuePosition::VECOUT, DOUBLE_BUFFER_NUM> outQueueRstd;185 TQue<QuePosition::VECOUT, DOUBLE_BUFFER_NUM> outQueueRstd;
186 TBuf<TPosition::VECCALC> xFp32Buf;186 TBuf<TPosition::VECCALC> xFp32Buf;
187 TBuf<TPosition::VECCALC> workLocalBuf;187 TBuf<TPosition::VECCALC> workLocalBuf;
188- MultiCopyParams<DX, NDDMA_DIM> dmaParam_;188+ NdDmaParams<DX, NDDMA_DIM> dmaParam_;
189 GlobalTensor<DX> xGm;189 GlobalTensor<DX> xGm;
190 GlobalTensor<DG> gammaGm;190 GlobalTensor<DG> gammaGm;
191 GlobalTensor<DX> yGm;191 GlobalTensor<DX> yGm;
Mnorm/rms_norm/op_kernel/arch35/rms_norm_regbase_common.h+170-170
@@ -154,7 +154,7 @@ __aicore__ inline void CopyOutY(GlobalTensor<T_Y>& yGm, TQue<QuePosition::VECOUT
154 */154 */
155__aicore__ inline void ComputeRstd(LocalTensor<float>& rstdLocal, float epsilon, float avgFactor, uint32_t count)155__aicore__ inline void ComputeRstd(LocalTensor<float>& rstdLocal, float epsilon, float avgFactor, uint32_t count)
156{156{
157- __local_mem__ float* rstdLocalAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();157+ __ubuf__ float* rstdLocalAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
158 158 
159 uint32_t calCount = count;159 uint32_t calCount = count;
160 uint32_t sreg = (uint32_t)calCount;160 uint32_t sreg = (uint32_t)calCount;
@@ -165,13 +165,13 @@ __aicore__ inline void ComputeRstd(LocalTensor<float>& rstdLocal, float epsilon,
165 MaskReg maskReg;165 MaskReg maskReg;
166 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {166 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
167 maskReg = UpdateMask<float>(sreg);167 maskReg = UpdateMask<float>(sreg);
168- DataCopy(srcReg, rstdLocalAddr + i * V_LENGTH);168+ LoadAlign(srcReg, rstdLocalAddr + i * V_LENGTH);
169 Muls(srcReg, srcReg, avgFactor, maskReg);169 Muls(srcReg, srcReg, avgFactor, maskReg);
170 Adds(dstReg, srcReg, epsilon, maskReg);170 Adds(dstReg, srcReg, epsilon, maskReg);
171 Sqrt(vReg, dstReg, maskReg);171 Sqrt(vReg, dstReg, maskReg);
172 Duplicate(srcReg, float(1.0), maskReg);172 Duplicate(srcReg, float(1.0), maskReg);
173 Div(dstReg, srcReg, vReg, maskReg);173 Div(dstReg, srcReg, vReg, maskReg);
174- DataCopy(rstdLocalAddr + i * V_LENGTH, dstReg, maskReg);174+ StoreAlign(rstdLocalAddr + i * V_LENGTH, dstReg, maskReg);
175 }175 }
176 }176 }
177}177}
@@ -194,18 +194,18 @@ __aicore__ inline void GemmaWithOutFloat(RegTensor<DG> gammaReg1, RegTensor<DG>
194}194}
195 195 
196template <typename DG, bool IS_GEMMA>196template <typename DG, bool IS_GEMMA>
197-__aicore__ inline void GemmaWithFloat(__local_mem__ DG* gammaAddr1, __local_mem__ DG* gammaAddr2,197+__aicore__ inline void GemmaWithFloat(__ubuf__ DG* gammaAddr1, __ubuf__ DG* gammaAddr2, RegTensor<DG>& gammaReg1,
198- RegTensor<DG>& gammaReg1, RegTensor<DG>& gammaReg2, MaskReg maskReg, uint16_t i)198+ RegTensor<DG>& gammaReg2, MaskReg maskReg, uint16_t i)
199{199{
200 if constexpr (IS_GEMMA) {200 if constexpr (IS_GEMMA) {
201 RegTensor<float> gammaTmp1, gammaTmp2;201 RegTensor<float> gammaTmp1, gammaTmp2;
202- DataCopy(gammaTmp1, gammaAddr1 + i * V_LENGTH);202+ LoadAlign(gammaTmp1, gammaAddr1 + i * V_LENGTH);
203- DataCopy(gammaTmp2, gammaAddr2 + i * V_LENGTH);203+ LoadAlign(gammaTmp2, gammaAddr2 + i * V_LENGTH);
204 Adds(gammaReg1, gammaTmp1, 1.0f, maskReg);204 Adds(gammaReg1, gammaTmp1, 1.0f, maskReg);
205 Adds(gammaReg2, gammaTmp2, 1.0f, maskReg);205 Adds(gammaReg2, gammaTmp2, 1.0f, maskReg);
206 } else {206 } else {
207- DataCopy(gammaReg1, gammaAddr1 + i * V_LENGTH);207+ LoadAlign(gammaReg1, gammaAddr1 + i * V_LENGTH);
208- DataCopy(gammaReg2, gammaAddr2 + i * V_LENGTH);208+ LoadAlign(gammaReg2, gammaAddr2 + i * V_LENGTH);
209 }209 }
210}210}
211 211 
@@ -226,13 +226,13 @@ __aicore__ inline void ComputeYMultiN(LocalTensor<float>& xLocal, LocalTensor<DG
226 uint32_t calCount = count / 2;226 uint32_t calCount = count / 2;
227 uint16_t repeatTimes = CeilDivision(calCount, V_LENGTH);227 uint16_t repeatTimes = CeilDivision(calCount, V_LENGTH);
228 228 
229- __local_mem__ float* xAddr1 = (__ubuf__ float*)xLocal.GetPhyAddr();229+ __ubuf__ float* xAddr1 = (__ubuf__ float*)xLocal.GetPhyAddr();
230- __local_mem__ float* xAddr2 = (__ubuf__ float*)xLocal.GetPhyAddr() + calCount;230+ __ubuf__ float* xAddr2 = (__ubuf__ float*)xLocal.GetPhyAddr() + calCount;
231- __local_mem__ DG* gammaAddr1 = (__ubuf__ DG*)gammaLocal.GetPhyAddr();231+ __ubuf__ DG* gammaAddr1 = (__ubuf__ DG*)gammaLocal.GetPhyAddr();
232- __local_mem__ DG* gammaAddr2 = (__ubuf__ DG*)gammaLocal.GetPhyAddr() + calCount;232+ __ubuf__ DG* gammaAddr2 = (__ubuf__ DG*)gammaLocal.GetPhyAddr() + calCount;
233- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();233+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
234- __local_mem__ DX* yAddr1 = (__ubuf__ DX*)yLocal.GetPhyAddr();234+ __ubuf__ DX* yAddr1 = (__ubuf__ DX*)yLocal.GetPhyAddr();
235- __local_mem__ DX* yAddr2 = (__ubuf__ DX*)yLocal.GetPhyAddr() + calCount;235+ __ubuf__ DX* yAddr2 = (__ubuf__ DX*)yLocal.GetPhyAddr() + calCount;
236 236 
237 if constexpr (!IsSameType<DX, float>::value && !IsSameType<DG, float>::value) {237 if constexpr (!IsSameType<DX, float>::value && !IsSameType<DG, float>::value) {
238 __VEC_SCOPE__238 __VEC_SCOPE__
@@ -245,13 +245,13 @@ __aicore__ inline void ComputeYMultiN(LocalTensor<float>& xLocal, LocalTensor<DG
245 RegTensor<float> xReg1, dst1Reg, gammaFp32Reg1, yReg1;245 RegTensor<float> xReg1, dst1Reg, gammaFp32Reg1, yReg1;
246 RegTensor<float> xReg2, dst2Reg, gammaFp32Reg2, yReg2;246 RegTensor<float> xReg2, dst2Reg, gammaFp32Reg2, yReg2;
247 MaskReg pregMask;247 MaskReg pregMask;
248- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + offset);248+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + offset);
249 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {249 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
250 pregMask = UpdateMask<float>(sreg);250 pregMask = UpdateMask<float>(sreg);
251- DataCopy(xReg1, xAddr1 + i * V_LENGTH);251+ LoadAlign(xReg1, xAddr1 + i * V_LENGTH);
252- DataCopy(xReg2, xAddr2 + i * V_LENGTH);252+ LoadAlign(xReg2, xAddr2 + i * V_LENGTH);
253- DataCopy<DG, LoadDist::DIST_UNPACK_B16>(gammaReg1, gammaAddr1 + i * V_LENGTH);253+ LoadAlign<DG, LoadDist::DIST_UNPACK_B16>(gammaReg1, gammaAddr1 + i * V_LENGTH);
254- DataCopy<DG, LoadDist::DIST_UNPACK_B16>(gammaReg2, gammaAddr2 + i * V_LENGTH);254+ LoadAlign<DG, LoadDist::DIST_UNPACK_B16>(gammaReg2, gammaAddr2 + i * V_LENGTH);
255 GemmaWithOutFloat<DG, IS_GEMMA>(gammaReg1, gammaReg2, gammaFp32Reg1, gammaFp32Reg2, pregMask);255 GemmaWithOutFloat<DG, IS_GEMMA>(gammaReg1, gammaReg2, gammaFp32Reg1, gammaFp32Reg2, pregMask);
256 Mul(dst1Reg, xReg1, rstdReg, pregMask);256 Mul(dst1Reg, xReg1, rstdReg, pregMask);
257 Mul(dst2Reg, xReg2, rstdReg, pregMask);257 Mul(dst2Reg, xReg2, rstdReg, pregMask);
@@ -259,8 +259,8 @@ __aicore__ inline void ComputeYMultiN(LocalTensor<float>& xLocal, LocalTensor<DG
259 Mul(yReg2, dst2Reg, gammaFp32Reg2, pregMask);259 Mul(yReg2, dst2Reg, gammaFp32Reg2, pregMask);
260 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, pregMask);260 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, pregMask);
261 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, pregMask);261 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, pregMask);
262- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, pregMask);262+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, pregMask);
263- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, pregMask);263+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, pregMask);
264 }264 }
265 offset++;265 offset++;
266 xAddr1 += count;266 xAddr1 += count;
@@ -280,11 +280,11 @@ __aicore__ inline void ComputeYMultiN(LocalTensor<float>& xLocal, LocalTensor<DG
280 RegTensor<float> xReg1, dst1Reg, yReg1;280 RegTensor<float> xReg1, dst1Reg, yReg1;
281 RegTensor<float> xReg2, dst2Reg, yReg2;281 RegTensor<float> xReg2, dst2Reg, yReg2;
282 MaskReg maskReg;282 MaskReg maskReg;
283- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + offset);283+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + offset);
284 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {284 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
285 maskReg = UpdateMask<float>(sreg);285 maskReg = UpdateMask<float>(sreg);
286- DataCopy(xReg1, xAddr1 + i * V_LENGTH);286+ LoadAlign(xReg1, xAddr1 + i * V_LENGTH);
287- DataCopy(xReg2, xAddr2 + i * V_LENGTH);287+ LoadAlign(xReg2, xAddr2 + i * V_LENGTH);
288 GemmaWithFloat<DG, IS_GEMMA>(gammaAddr1, gammaAddr2, gammaReg1, gammaReg2, maskReg, i);288 GemmaWithFloat<DG, IS_GEMMA>(gammaAddr1, gammaAddr2, gammaReg1, gammaReg2, maskReg, i);
289 Mul(dst1Reg, xReg1, rstdReg, maskReg);289 Mul(dst1Reg, xReg1, rstdReg, maskReg);
290 Mul(dst2Reg, xReg2, rstdReg, maskReg);290 Mul(dst2Reg, xReg2, rstdReg, maskReg);
@@ -292,8 +292,8 @@ __aicore__ inline void ComputeYMultiN(LocalTensor<float>& xLocal, LocalTensor<DG
292 Mul(yReg2, dst2Reg, gammaReg2, maskReg);292 Mul(yReg2, dst2Reg, gammaReg2, maskReg);
293 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, maskReg);293 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, maskReg);
294 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, maskReg);294 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, maskReg);
295- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, maskReg);295+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, maskReg);
296- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, maskReg);296+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, maskReg);
297 }297 }
298 offset++;298 offset++;
299 xAddr1 += count;299 xAddr1 += count;
@@ -311,18 +311,18 @@ __aicore__ inline void ComputeYMultiN(LocalTensor<float>& xLocal, LocalTensor<DG
311 RegTensor<float> xReg1, gammaReg1, yReg1, vRegTmp1;311 RegTensor<float> xReg1, gammaReg1, yReg1, vRegTmp1;
312 RegTensor<float> xReg2, gammaReg2, yReg2, vRegTmp2;312 RegTensor<float> xReg2, gammaReg2, yReg2, vRegTmp2;
313 MaskReg maskReg;313 MaskReg maskReg;
314- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + offset);314+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + offset);
315 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {315 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
316 maskReg = UpdateMask<float>(sreg);316 maskReg = UpdateMask<float>(sreg);
317- DataCopy(xReg1, xAddr1 + i * V_LENGTH);317+ LoadAlign(xReg1, xAddr1 + i * V_LENGTH);
318- DataCopy(xReg2, xAddr2 + i * V_LENGTH);318+ LoadAlign(xReg2, xAddr2 + i * V_LENGTH);
319 GemmaWithFloat<DG, IS_GEMMA>(gammaAddr1, gammaAddr2, gammaReg1, gammaReg2, maskReg, i);319 GemmaWithFloat<DG, IS_GEMMA>(gammaAddr1, gammaAddr2, gammaReg1, gammaReg2, maskReg, i);
320 Mul(vRegTmp1, xReg1, rstdReg, maskReg);320 Mul(vRegTmp1, xReg1, rstdReg, maskReg);
321 Mul(vRegTmp2, xReg2, rstdReg, maskReg);321 Mul(vRegTmp2, xReg2, rstdReg, maskReg);
322 Mul(yReg1, vRegTmp1, gammaReg1, maskReg);322 Mul(yReg1, vRegTmp1, gammaReg1, maskReg);
323 Mul(yReg2, vRegTmp2, gammaReg2, maskReg);323 Mul(yReg2, vRegTmp2, gammaReg2, maskReg);
324- DataCopy(yAddr1 + i * V_LENGTH, yReg1, maskReg);324+ StoreAlign(yAddr1 + i * V_LENGTH, yReg1, maskReg);
325- DataCopy(yAddr2 + i * V_LENGTH, yReg2, maskReg);325+ StoreAlign(yAddr2 + i * V_LENGTH, yReg2, maskReg);
326 }326 }
327 offset++;327 offset++;
328 xAddr1 += count;328 xAddr1 += count;
@@ -352,13 +352,13 @@ __aicore__ inline void ComputeLatterY(LocalTensor<DX>& xLocal, LocalTensor<DG>&
352 uint32_t sreg = (uint32_t)calCount;352 uint32_t sreg = (uint32_t)calCount;
353 uint16_t repeatTimes = CeilDivision(calCount, V_LENGTH);353 uint16_t repeatTimes = CeilDivision(calCount, V_LENGTH);
354 354 
355- __local_mem__ DX* xAddr1 = (__ubuf__ DX*)xLocal.GetPhyAddr();355+ __ubuf__ DX* xAddr1 = (__ubuf__ DX*)xLocal.GetPhyAddr();
356- __local_mem__ DX* xAddr2 = (__ubuf__ DX*)xLocal.GetPhyAddr() + calCount;356+ __ubuf__ DX* xAddr2 = (__ubuf__ DX*)xLocal.GetPhyAddr() + calCount;
357- __local_mem__ DG* gammaAddr1 = (__ubuf__ DG*)gammaLocal.GetPhyAddr();357+ __ubuf__ DG* gammaAddr1 = (__ubuf__ DG*)gammaLocal.GetPhyAddr();
358- __local_mem__ DG* gammaAddr2 = (__ubuf__ DG*)gammaLocal.GetPhyAddr() + calCount;358+ __ubuf__ DG* gammaAddr2 = (__ubuf__ DG*)gammaLocal.GetPhyAddr() + calCount;
359- __local_mem__ float* srcAddr2 = (__ubuf__ float*)rstdLocal.GetPhyAddr();359+ __ubuf__ float* srcAddr2 = (__ubuf__ float*)rstdLocal.GetPhyAddr();
360- __local_mem__ DX* yAddr1 = (__ubuf__ DX*)yLocal.GetPhyAddr();360+ __ubuf__ DX* yAddr1 = (__ubuf__ DX*)yLocal.GetPhyAddr();
361- __local_mem__ DX* yAddr2 = (__ubuf__ DX*)yLocal.GetPhyAddr() + calCount;361+ __ubuf__ DX* yAddr2 = (__ubuf__ DX*)yLocal.GetPhyAddr() + calCount;
362 362 
363 if constexpr (!IsSameType<DX, float>::value and !IsSameType<DG, float>::value) {363 if constexpr (!IsSameType<DX, float>::value and !IsSameType<DG, float>::value) {
364 __VEC_SCOPE__364 __VEC_SCOPE__
@@ -369,13 +369,13 @@ __aicore__ inline void ComputeLatterY(LocalTensor<DX>& xLocal, LocalTensor<DG>&
369 RegTensor<float> xReg1, dst1Reg, gammaFp32Reg1, yReg1;369 RegTensor<float> xReg1, dst1Reg, gammaFp32Reg1, yReg1;
370 RegTensor<float> xReg2, dst2Reg, gammaFp32Reg2, yReg2;370 RegTensor<float> xReg2, dst2Reg, gammaFp32Reg2, yReg2;
371 MaskReg maskReg;371 MaskReg maskReg;
372- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, srcAddr2 + offset);372+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, srcAddr2 + offset);
373 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {373 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
374 maskReg = UpdateMask<float>(sreg);374 maskReg = UpdateMask<float>(sreg);
375- DataCopy<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg1, xAddr1 + i * V_LENGTH);375+ LoadAlign<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg1, xAddr1 + i * V_LENGTH);
376- DataCopy<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg2, xAddr2 + i * V_LENGTH);376+ LoadAlign<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg2, xAddr2 + i * V_LENGTH);
377- DataCopy<DG, LoadDist::DIST_UNPACK_B16>(gammaReg1, gammaAddr1 + i * V_LENGTH);377+ LoadAlign<DG, LoadDist::DIST_UNPACK_B16>(gammaReg1, gammaAddr1 + i * V_LENGTH);
378- DataCopy<DG, LoadDist::DIST_UNPACK_B16>(gammaReg2, gammaAddr2 + i * V_LENGTH);378+ LoadAlign<DG, LoadDist::DIST_UNPACK_B16>(gammaReg2, gammaAddr2 + i * V_LENGTH);
379 if constexpr (IS_GEMMA) {379 if constexpr (IS_GEMMA) {
380 RegTensor<float> gammaTmp1, gammaTmp2;380 RegTensor<float> gammaTmp1, gammaTmp2;
381 Cast<float, DG, castTraitB162B32>(gammaTmp1, gammaReg1, maskReg);381 Cast<float, DG, castTraitB162B32>(gammaTmp1, gammaReg1, maskReg);
@@ -394,8 +394,8 @@ __aicore__ inline void ComputeLatterY(LocalTensor<DX>& xLocal, LocalTensor<DG>&
394 Mul(yReg2, dst2Reg, gammaFp32Reg2, maskReg);394 Mul(yReg2, dst2Reg, gammaFp32Reg2, maskReg);
395 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, maskReg);395 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, maskReg);
396 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, maskReg);396 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, maskReg);
397- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, maskReg);397+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, maskReg);
398- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, maskReg);398+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, maskReg);
399 }399 }
400 }400 }
401 } else if constexpr (!IsSameType<DX, float>::value and IsSameType<DG, float>::value) {401 } else if constexpr (!IsSameType<DX, float>::value and IsSameType<DG, float>::value) {
@@ -406,20 +406,20 @@ __aicore__ inline void ComputeLatterY(LocalTensor<DX>& xLocal, LocalTensor<DG>&
406 RegTensor<float> xReg1, dst1Reg, gammaFp32Reg1, yReg1;406 RegTensor<float> xReg1, dst1Reg, gammaFp32Reg1, yReg1;
407 RegTensor<float> xReg2, dst2Reg, gammaFp32Reg2, yReg2;407 RegTensor<float> xReg2, dst2Reg, gammaFp32Reg2, yReg2;
408 MaskReg maskReg;408 MaskReg maskReg;
409- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, srcAddr2 + offset);409+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, srcAddr2 + offset);
410 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {410 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
411 maskReg = UpdateMask<float>(sreg);411 maskReg = UpdateMask<float>(sreg);
412- DataCopy<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg1, xAddr1 + i * V_LENGTH);412+ LoadAlign<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg1, xAddr1 + i * V_LENGTH);
413- DataCopy<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg2, xAddr2 + i * V_LENGTH);413+ LoadAlign<DX, LoadDist::DIST_UNPACK_B16>(xB16Reg2, xAddr2 + i * V_LENGTH);
414 if constexpr (IS_GEMMA) {414 if constexpr (IS_GEMMA) {
415 RegTensor<float> gammaTmp1, gammaTmp2;415 RegTensor<float> gammaTmp1, gammaTmp2;
416- DataCopy(gammaTmp1, gammaAddr1 + i * V_LENGTH);416+ LoadAlign(gammaTmp1, gammaAddr1 + i * V_LENGTH);
417- DataCopy(gammaTmp2, gammaAddr2 + i * V_LENGTH);417+ LoadAlign(gammaTmp2, gammaAddr2 + i * V_LENGTH);
418 Adds(gammaFp32Reg1, gammaTmp1, 1.0f, maskReg);418 Adds(gammaFp32Reg1, gammaTmp1, 1.0f, maskReg);
419 Adds(gammaFp32Reg2, gammaTmp2, 1.0f, maskReg);419 Adds(gammaFp32Reg2, gammaTmp2, 1.0f, maskReg);
420 } else {420 } else {
421- DataCopy(gammaFp32Reg1, gammaAddr1 + i * V_LENGTH);421+ LoadAlign(gammaFp32Reg1, gammaAddr1 + i * V_LENGTH);
422- DataCopy(gammaFp32Reg2, gammaAddr2 + i * V_LENGTH);422+ LoadAlign(gammaFp32Reg2, gammaAddr2 + i * V_LENGTH);
423 }423 }
424 Cast<float, DX, castTraitB162B32>(xReg1, xB16Reg1, maskReg);424 Cast<float, DX, castTraitB162B32>(xReg1, xB16Reg1, maskReg);
425 Cast<float, DX, castTraitB162B32>(xReg2, xB16Reg2, maskReg);425 Cast<float, DX, castTraitB162B32>(xReg2, xB16Reg2, maskReg);
@@ -429,8 +429,8 @@ __aicore__ inline void ComputeLatterY(LocalTensor<DX>& xLocal, LocalTensor<DG>&
429 Mul(yReg2, dst2Reg, gammaFp32Reg2, maskReg);429 Mul(yReg2, dst2Reg, gammaFp32Reg2, maskReg);
430 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, maskReg);430 Cast<DX, float, castTraitB322B16>(yB16Reg1, yReg1, maskReg);
431 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, maskReg);431 Cast<DX, float, castTraitB322B16>(yB16Reg2, yReg2, maskReg);
432- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, maskReg);432+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr1 + i * V_LENGTH, yB16Reg1, maskReg);
433- DataCopy<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, maskReg);433+ StoreAlign<DX, StoreDist::DIST_PACK_B32>(yAddr2 + i * V_LENGTH, yB16Reg2, maskReg);
434 }434 }
435 }435 }
436 } else {436 } else {
@@ -440,27 +440,27 @@ __aicore__ inline void ComputeLatterY(LocalTensor<DX>& xLocal, LocalTensor<DG>&
440 RegTensor<float> xReg1, gammaReg1, yReg1, vRegTmp1;440 RegTensor<float> xReg1, gammaReg1, yReg1, vRegTmp1;
441 RegTensor<float> xReg2, gammaReg2, yReg2, vRegTmp2;441 RegTensor<float> xReg2, gammaReg2, yReg2, vRegTmp2;
442 MaskReg maskReg;442 MaskReg maskReg;
443- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, srcAddr2 + offset);443+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, srcAddr2 + offset);
444 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {444 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
445 maskReg = UpdateMask<float>(sreg);445 maskReg = UpdateMask<float>(sreg);
446- DataCopy(xReg1, xAddr1 + i * V_LENGTH);446+ LoadAlign(xReg1, xAddr1 + i * V_LENGTH);
447- DataCopy(xReg2, xAddr2 + i * V_LENGTH);447+ LoadAlign(xReg2, xAddr2 + i * V_LENGTH);
448 if constexpr (IS_GEMMA) {448 if constexpr (IS_GEMMA) {
449 RegTensor<float> gammaTmp1, gammaTmp2;449 RegTensor<float> gammaTmp1, gammaTmp2;
450- DataCopy(gammaTmp1, gammaAddr1 + i * V_LENGTH);450+ LoadAlign(gammaTmp1, gammaAddr1 + i * V_LENGTH);
451- DataCopy(gammaTmp2, gammaAddr2 + i * V_LENGTH);451+ LoadAlign(gammaTmp2, gammaAddr2 + i * V_LENGTH);
452 Adds(gammaReg1, gammaTmp1, 1.0f, maskReg);452 Adds(gammaReg1, gammaTmp1, 1.0f, maskReg);
453 Adds(gammaReg2, gammaTmp2, 1.0f, maskReg);453 Adds(gammaReg2, gammaTmp2, 1.0f, maskReg);
454 } else {454 } else {
455- DataCopy(gammaReg1, gammaAddr1 + i * V_LENGTH);455+ LoadAlign(gammaReg1, gammaAddr1 + i * V_LENGTH);
456- DataCopy(gammaReg2, gammaAddr2 + i * V_LENGTH);456+ LoadAlign(gammaReg2, gammaAddr2 + i * V_LENGTH);
457 }457 }
458 Mul(vRegTmp1, xReg1, rstdReg, maskReg);458 Mul(vRegTmp1, xReg1, rstdReg, maskReg);
459 Mul(vRegTmp2, xReg2, rstdReg, maskReg);459 Mul(vRegTmp2, xReg2, rstdReg, maskReg);
460 Mul(yReg1, vRegTmp1, gammaReg1, maskReg);460 Mul(yReg1, vRegTmp1, gammaReg1, maskReg);
461 Mul(yReg2, vRegTmp2, gammaReg2, maskReg);461 Mul(yReg2, vRegTmp2, gammaReg2, maskReg);
462- DataCopy(yAddr1 + i * V_LENGTH, yReg1, maskReg);462+ StoreAlign(yAddr1 + i * V_LENGTH, yReg1, maskReg);
463- DataCopy(yAddr2 + i * V_LENGTH, yReg2, maskReg);463+ StoreAlign(yAddr2 + i * V_LENGTH, yReg2, maskReg);
464 }464 }
465 }465 }
466 }466 }
@@ -498,8 +498,8 @@ __aicore__ inline void ComputeSum(LocalTensor<float>& dstLocal, LocalTensor<floa
498 uint32_t meanTile = count;498 uint32_t meanTile = count;
499 uint32_t meanSreg = meanTile;499 uint32_t meanSreg = meanTile;
500 500 
501- __local_mem__ float* srcAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();501+ __ubuf__ float* srcAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();
502- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();502+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
503 503 
504 __VEC_SCOPE__504 __VEC_SCOPE__
505 {505 {
@@ -508,57 +508,57 @@ __aicore__ inline void ComputeSum(LocalTensor<float>& dstLocal, LocalTensor<floa
508 MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();508 MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();
509 {509 {
510 pregLoop = UpdateMask<float>(meanSreg);510 pregLoop = UpdateMask<float>(meanSreg);
511- DataCopy(vReg, srcAddr + 0);511+ LoadAlign(vReg, srcAddr + 0);
512- ReduceSum(vMean, vReg, pregLoop);512+ Reduce<ReduceType::SUM>(vMean, vReg, pregLoop);
513- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMean, pregMerge);513+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMean, pregMerge);
514 }514 }
515 }515 }
516}516}
517 517 
518template <typename T, bool SAVE_FP32>518template <typename T, bool SAVE_FP32>
519-__aicore__ inline void LoadSquareRemainTile(__local_mem__ T* mainAddr, __local_mem__ T* tailAddr, uint16_t offset1,519+__aicore__ inline void LoadSquareRemainTile(__ubuf__ T* mainAddr, __ubuf__ T* tailAddr, uint16_t offset1,
520 uint16_t offset2, RegTensor<float>& mainA, RegTensor<float>& mainB,520 uint16_t offset2, RegTensor<float>& mainA, RegTensor<float>& mainB,
521 RegTensor<float>& tailA, RegTensor<float>& tailB, MaskReg& pregLoop,521 RegTensor<float>& tailA, RegTensor<float>& tailB, MaskReg& pregLoop,
522- __local_mem__ float* xFp32MainAddr = nullptr,522+ __ubuf__ float* xFp32MainAddr = nullptr,
523- __local_mem__ float* xFp32TailAddr = nullptr)523+ __ubuf__ float* xFp32TailAddr = nullptr)
524{524{
525 if constexpr (IsSameType<T, half>::value) {525 if constexpr (IsSameType<T, half>::value) {
526 RegTensor<half> xFp16MainA, xFp16MainB, xFp16TailA, xFp16TailB;526 RegTensor<half> xFp16MainA, xFp16MainB, xFp16TailA, xFp16TailB;
527- DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16MainA, mainAddr + offset1);527+ LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16MainA, mainAddr + offset1);
528- DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16MainB, mainAddr + offset2);528+ LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16MainB, mainAddr + offset2);
529- DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16TailA, tailAddr + offset1);529+ LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16TailA, tailAddr + offset1);
530- DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16TailB, tailAddr + offset2);530+ LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16TailB, tailAddr + offset2);
531 Cast<float, half, castTraitB162B32>(mainA, xFp16MainA, pregLoop);531 Cast<float, half, castTraitB162B32>(mainA, xFp16MainA, pregLoop);
532 Cast<float, half, castTraitB162B32>(mainB, xFp16MainB, pregLoop);532 Cast<float, half, castTraitB162B32>(mainB, xFp16MainB, pregLoop);
533 Cast<float, half, castTraitB162B32>(tailA, xFp16TailA, pregLoop);533 Cast<float, half, castTraitB162B32>(tailA, xFp16TailA, pregLoop);
534 Cast<float, half, castTraitB162B32>(tailB, xFp16TailB, pregLoop);534 Cast<float, half, castTraitB162B32>(tailB, xFp16TailB, pregLoop);
535 if constexpr (SAVE_FP32) {535 if constexpr (SAVE_FP32) {
536- DataCopy(xFp32MainAddr + offset1, mainA, pregLoop);536+ StoreAlign(xFp32MainAddr + offset1, mainA, pregLoop);
537- DataCopy(xFp32MainAddr + offset2, mainB, pregLoop);537+ StoreAlign(xFp32MainAddr + offset2, mainB, pregLoop);
538- DataCopy(xFp32TailAddr + offset1, tailA, pregLoop);538+ StoreAlign(xFp32TailAddr + offset1, tailA, pregLoop);
539- DataCopy(xFp32TailAddr + offset2, tailB, pregLoop);539+ StoreAlign(xFp32TailAddr + offset2, tailB, pregLoop);
540 }540 }
541 } else if constexpr (IsSameType<T, bfloat16_t>::value) {541 } else if constexpr (IsSameType<T, bfloat16_t>::value) {
542 RegTensor<bfloat16_t> xBFp16MainA, xBFp16MainB, xBFp16TailA, xBFp16TailB;542 RegTensor<bfloat16_t> xBFp16MainA, xBFp16MainB, xBFp16TailA, xBFp16TailB;
543- DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainA, mainAddr + offset1);543+ LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainA, mainAddr + offset1);
544- DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainB, mainAddr + offset2);544+ LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainB, mainAddr + offset2);
545- DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16TailA, tailAddr + offset1);545+ LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16TailA, tailAddr + offset1);
546- DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16TailB, tailAddr + offset2);546+ LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16TailB, tailAddr + offset2);
547 Cast<float, bfloat16_t, castTraitB162B32>(mainA, xBFp16MainA, pregLoop);547 Cast<float, bfloat16_t, castTraitB162B32>(mainA, xBFp16MainA, pregLoop);
548 Cast<float, bfloat16_t, castTraitB162B32>(mainB, xBFp16MainB, pregLoop);548 Cast<float, bfloat16_t, castTraitB162B32>(mainB, xBFp16MainB, pregLoop);
549 Cast<float, bfloat16_t, castTraitB162B32>(tailA, xBFp16TailA, pregLoop);549 Cast<float, bfloat16_t, castTraitB162B32>(tailA, xBFp16TailA, pregLoop);
550 Cast<float, bfloat16_t, castTraitB162B32>(tailB, xBFp16TailB, pregLoop);550 Cast<float, bfloat16_t, castTraitB162B32>(tailB, xBFp16TailB, pregLoop);
551 if constexpr (SAVE_FP32) {551 if constexpr (SAVE_FP32) {
552- DataCopy(xFp32MainAddr + offset1, mainA, pregLoop);552+ StoreAlign(xFp32MainAddr + offset1, mainA, pregLoop);
553- DataCopy(xFp32MainAddr + offset2, mainB, pregLoop);553+ StoreAlign(xFp32MainAddr + offset2, mainB, pregLoop);
554- DataCopy(xFp32TailAddr + offset1, tailA, pregLoop);554+ StoreAlign(xFp32TailAddr + offset1, tailA, pregLoop);
555- DataCopy(xFp32TailAddr + offset2, tailB, pregLoop);555+ StoreAlign(xFp32TailAddr + offset2, tailB, pregLoop);
556 }556 }
557 } else {557 } else {
558- DataCopy(mainA, mainAddr + offset1);558+ LoadAlign(mainA, mainAddr + offset1);
559- DataCopy(mainB, mainAddr + offset2);559+ LoadAlign(mainB, mainAddr + offset2);
560- DataCopy(tailA, tailAddr + offset1);560+ LoadAlign(tailA, tailAddr + offset1);
561- DataCopy(tailB, tailAddr + offset2);561+ LoadAlign(tailB, tailAddr + offset2);
562 }562 }
563 Mul(mainA, mainA, mainA, pregLoop);563 Mul(mainA, mainA, mainA, pregLoop);
564 Mul(mainB, mainB, mainB, pregLoop);564 Mul(mainB, mainB, mainB, pregLoop);
@@ -567,37 +567,37 @@ __aicore__ inline void LoadSquareRemainTile(__local_mem__ T* mainAddr, __local_m
567}567}
568 568 
569template <typename T, bool SAVE_FP32>569template <typename T, bool SAVE_FP32>
570-__aicore__ inline void LoadSquareMasterTile(__local_mem__ T* masterAddr, uint16_t offset1, uint16_t offset2,570+__aicore__ inline void LoadSquareMasterTile(__ubuf__ T* masterAddr, uint16_t offset1, uint16_t offset2,
571 RegTensor<float>& mainA, RegTensor<float>& mainB, MaskReg& pregLoop,571 RegTensor<float>& mainA, RegTensor<float>& mainB, MaskReg& pregLoop,
572- __local_mem__ float* xFp32MasterAddr = nullptr)572+ __ubuf__ float* xFp32MasterAddr = nullptr)
573{573{
574 if constexpr (IsSameType<T, half>::value) {574 if constexpr (IsSameType<T, half>::value) {
575 RegTensor<half> xFp16MainA, xFp16MainB;575 RegTensor<half> xFp16MainA, xFp16MainB;
576- DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16MainA, masterAddr + offset1);576+ LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16MainA, masterAddr + offset1);
577- DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16MainB, masterAddr + offset2);577+ LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16MainB, masterAddr + offset2);
578 Cast<float, half, castTraitB162B32>(mainA, xFp16MainA, pregLoop);578 Cast<float, half, castTraitB162B32>(mainA, xFp16MainA, pregLoop);
579 Cast<float, half, castTraitB162B32>(mainB, xFp16MainB, pregLoop);579 Cast<float, half, castTraitB162B32>(mainB, xFp16MainB, pregLoop);
580 if constexpr (SAVE_FP32) {580 if constexpr (SAVE_FP32) {
581- DataCopy(xFp32MasterAddr + offset1, mainA, pregLoop);581+ StoreAlign(xFp32MasterAddr + offset1, mainA, pregLoop);
582- DataCopy(xFp32MasterAddr + offset2, mainB, pregLoop);582+ StoreAlign(xFp32MasterAddr + offset2, mainB, pregLoop);
583 }583 }
584 Mul(mainA, mainA, mainA, pregLoop);584 Mul(mainA, mainA, mainA, pregLoop);
585 Mul(mainB, mainB, mainB, pregLoop);585 Mul(mainB, mainB, mainB, pregLoop);
586 } else if constexpr (IsSameType<T, bfloat16_t>::value) {586 } else if constexpr (IsSameType<T, bfloat16_t>::value) {
587 RegTensor<bfloat16_t> xBFp16MainA, xBFp16MainB;587 RegTensor<bfloat16_t> xBFp16MainA, xBFp16MainB;
588- DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainA, masterAddr + offset1);588+ LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainA, masterAddr + offset1);
589- DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainB, masterAddr + offset2);589+ LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBFp16MainB, masterAddr + offset2);
590 Cast<float, bfloat16_t, castTraitB162B32>(mainA, xBFp16MainA, pregLoop);590 Cast<float, bfloat16_t, castTraitB162B32>(mainA, xBFp16MainA, pregLoop);
591 Cast<float, bfloat16_t, castTraitB162B32>(mainB, xBFp16MainB, pregLoop);591 Cast<float, bfloat16_t, castTraitB162B32>(mainB, xBFp16MainB, pregLoop);
592 if constexpr (SAVE_FP32) {592 if constexpr (SAVE_FP32) {
593- DataCopy(xFp32MasterAddr + offset1, mainA, pregLoop);593+ StoreAlign(xFp32MasterAddr + offset1, mainA, pregLoop);
594- DataCopy(xFp32MasterAddr + offset2, mainB, pregLoop);594+ StoreAlign(xFp32MasterAddr + offset2, mainB, pregLoop);
595 }595 }
596 Mul(mainA, mainA, mainA, pregLoop);596 Mul(mainA, mainA, mainA, pregLoop);
597 Mul(mainB, mainB, mainB, pregLoop);597 Mul(mainB, mainB, mainB, pregLoop);
598 } else {598 } else {
599- DataCopy(mainA, masterAddr + offset1);599+ LoadAlign(mainA, masterAddr + offset1);
600- DataCopy(mainB, masterAddr + offset2);600+ LoadAlign(mainB, masterAddr + offset2);
601 Mul(mainA, mainA, mainA, pregLoop);601 Mul(mainA, mainA, mainA, pregLoop);
602 Mul(mainB, mainB, mainB, pregLoop);602 Mul(mainB, mainB, mainB, pregLoop);
603 }603 }
@@ -620,10 +620,10 @@ __aicore__ inline void ComputeFormerImplV1MultiN(LocalTensor<T>& xLocal, LocalTe
620 620 
621 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;621 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;
622 622 
623- __local_mem__ T* mainAddr = (__ubuf__ T*)xLocal.GetPhyAddr();623+ __ubuf__ T* mainAddr = (__ubuf__ T*)xLocal.GetPhyAddr();
624- __local_mem__ T* tailAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit);624+ __ubuf__ T* tailAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit);
625- __local_mem__ T* masterAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile);625+ __ubuf__ T* masterAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile);
626- __local_mem__ float *xFp32MainAddr, *xFp32TailAddr, *xFp32MasterAddr;626+ __ubuf__ float *xFp32MainAddr, *xFp32TailAddr, *xFp32MasterAddr;
627 if constexpr (is_same<T, half>::value || is_same<T, bfloat16_t>::value) {627 if constexpr (is_same<T, half>::value || is_same<T, bfloat16_t>::value) {
628 xFp32MainAddr = (__ubuf__ float*)xFp32.GetPhyAddr();628 xFp32MainAddr = (__ubuf__ float*)xFp32.GetPhyAddr();
629 xFp32TailAddr = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(powerSplit);629 xFp32TailAddr = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(powerSplit);
@@ -634,21 +634,21 @@ __aicore__ inline void ComputeFormerImplV1MultiN(LocalTensor<T>& xLocal, LocalTe
634 bool isWithTail = curRowsAlign - (curRows / 2);634 bool isWithTail = curRowsAlign - (curRows / 2);
635 uint32_t tailOffset = offset + curRows / 2;635 uint32_t tailOffset = offset + curRows / 2;
636 636 
637- __local_mem__ T* mainAddr1 = (__ubuf__ T*)xLocal.GetPhyAddr() + unrollOffset;637+ __ubuf__ T* mainAddr1 = (__ubuf__ T*)xLocal.GetPhyAddr() + unrollOffset;
638- __local_mem__ T* tailAddr1 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit) + unrollOffset;638+ __ubuf__ T* tailAddr1 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit) + unrollOffset;
639- __local_mem__ T* masterAddr1 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile) + unrollOffset;639+ __ubuf__ T* masterAddr1 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile) + unrollOffset;
640- __local_mem__ float *xFp32MainAddr1, *xFp32TailAddr1, *xFp32MasterAddr1;640+ __ubuf__ float *xFp32MainAddr1, *xFp32TailAddr1, *xFp32MasterAddr1;
641 if constexpr (is_same<T, half>::value || is_same<T, bfloat16_t>::value) {641 if constexpr (is_same<T, half>::value || is_same<T, bfloat16_t>::value) {
642 xFp32MainAddr1 = (__ubuf__ float*)xFp32.GetPhyAddr() + unrollOffset;642 xFp32MainAddr1 = (__ubuf__ float*)xFp32.GetPhyAddr() + unrollOffset;
643 xFp32TailAddr1 = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(powerSplit) + unrollOffset;643 xFp32TailAddr1 = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(powerSplit) + unrollOffset;
644 xFp32MasterAddr1 = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(remainTile) + unrollOffset;644 xFp32MasterAddr1 = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(remainTile) + unrollOffset;
645 }645 }
646 646 
647- __local_mem__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();647+ __ubuf__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();
648- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();648+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
649 649 
650- __local_mem__ float* workAddr1 = (__ubuf__ float*)workLocal.GetPhyAddr() + NormCommon::ONCE_VECTOR_SIZE;650+ __ubuf__ float* workAddr1 = (__ubuf__ float*)workLocal.GetPhyAddr() + NormCommon::ONCE_VECTOR_SIZE;
651- __local_mem__ float* rstdAddr1 = (__ubuf__ float*)rstdLocal.GetPhyAddr() + curRows / 2;651+ __ubuf__ float* rstdAddr1 = (__ubuf__ float*)rstdLocal.GetPhyAddr() + curRows / 2;
652 652 
653 __VEC_SCOPE__653 __VEC_SCOPE__
654 {654 {
@@ -677,16 +677,16 @@ __aicore__ inline void ComputeFormerImplV1MultiN(LocalTensor<T>& xLocal, LocalTe
677 Add(mainA, mainA, tailA, pregLoop);677 Add(mainA, mainA, tailA, pregLoop);
678 Add(mainB, mainB, tailB, pregLoop);678 Add(mainB, mainB, tailB, pregLoop);
679 Add(mainA, mainA, mainB, pregLoop);679 Add(mainA, mainA, mainB, pregLoop);
680- ReduceSum(vMean, mainA, pregLoop);680+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
681- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);681+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);
682 }682 }
683 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {683 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
684 pregLoop = UpdateMask<float>(masterSreg);684 pregLoop = UpdateMask<float>(masterSreg);
685 LoadSquareMasterTile<T, true>(masterAddr, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA, mainB,685 LoadSquareMasterTile<T, true>(masterAddr, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA, mainB,
686 pregLoop, xFp32MasterAddr);686 pregLoop, xFp32MasterAddr);
687 Add(mainA, mainA, mainB, pregLoop);687 Add(mainA, mainA, mainB, pregLoop);
688- ReduceSum(vMean, mainA, pregLoop);688+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
689- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + remainRepeats + i, vMean, pregMerge);689+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + remainRepeats + i, vMean, pregMerge);
690 }690 }
691 // unroll part691 // unroll part
692 for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {692 for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {
@@ -697,58 +697,58 @@ __aicore__ inline void ComputeFormerImplV1MultiN(LocalTensor<T>& xLocal, LocalTe
697 Add(mainA1, mainA1, tailA1, pregLoop1);697 Add(mainA1, mainA1, tailA1, pregLoop1);
698 Add(mainB1, mainB1, tailB1, pregLoop1);698 Add(mainB1, mainB1, tailB1, pregLoop1);
699 Add(mainA1, mainA1, mainB1, pregLoop1);699 Add(mainA1, mainA1, mainB1, pregLoop1);
700- ReduceSum(vMean1, mainA1, pregLoop1);700+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
701- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);701+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);
702 }702 }
703 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {703 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
704 pregLoop1 = UpdateMask<float>(masterSreg1);704 pregLoop1 = UpdateMask<float>(masterSreg1);
705 LoadSquareMasterTile<T, true>(masterAddr1, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA1,705 LoadSquareMasterTile<T, true>(masterAddr1, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA1,
706 mainB1, pregLoop1, xFp32MasterAddr1);706 mainB1, pregLoop1, xFp32MasterAddr1);
707 Add(mainA1, mainA1, mainB1, pregLoop1);707 Add(mainA1, mainA1, mainB1, pregLoop1);
708- ReduceSum(vMean1, mainA1, pregLoop1);708+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
709- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + remainRepeats + i, vMean1, pregMerge1);709+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + remainRepeats + i, vMean1, pregMerge1);
710 }710 }
711 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();711 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
712 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {712 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
713 pregLoop = UpdateMask<float>(mergeSreg);713 pregLoop = UpdateMask<float>(mergeSreg);
714- DataCopy(mainA, workAddr + (i * 2 + 0) * V_LENGTH);714+ LoadAlign(mainA, workAddr + (i * 2 + 0) * V_LENGTH);
715- DataCopy(mainB, workAddr + (i * 2 + 1) * V_LENGTH);715+ LoadAlign(mainB, workAddr + (i * 2 + 1) * V_LENGTH);
716 Add(mainA, mainA, mainB, pregLoop);716 Add(mainA, mainA, mainB, pregLoop);
717- ReduceSum(vMean, mainA, pregLoop);717+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
718- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);718+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);
719 }719 }
720 // unroll part720 // unroll part
721 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {721 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
722 pregLoop1 = UpdateMask<float>(mergeSreg1);722 pregLoop1 = UpdateMask<float>(mergeSreg1);
723- DataCopy(mainA1, workAddr1 + (i * 2 + 0) * V_LENGTH);723+ LoadAlign(mainA1, workAddr1 + (i * 2 + 0) * V_LENGTH);
724- DataCopy(mainB1, workAddr1 + (i * 2 + 1) * V_LENGTH);724+ LoadAlign(mainB1, workAddr1 + (i * 2 + 1) * V_LENGTH);
725 Add(mainA1, mainA1, mainB1, pregLoop1);725 Add(mainA1, mainA1, mainB1, pregLoop1);
726- ReduceSum(vMean1, mainA1, pregLoop1);726+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
727- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);727+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);
728 }728 }
729 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();729 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
730 {730 {
731 pregLoop = UpdateMask<float>(meanSreg);731 pregLoop = UpdateMask<float>(meanSreg);
732- DataCopy(mainA, workAddr + 0);732+ LoadAlign(mainA, workAddr + 0);
733- ReduceSum(vMean, mainA, pregLoop);733+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
734 Muls(vMean, vMean, avgFactor, pregMerge);734 Muls(vMean, vMean, avgFactor, pregMerge);
735 Adds(vMean, vMean, epsilon, pregMerge);735 Adds(vMean, vMean, epsilon, pregMerge);
736 Sqrt(vMean, vMean, pregMerge);736 Sqrt(vMean, vMean, pregMerge);
737 Duplicate(vDupReg, float(1.0), pregMerge);737 Duplicate(vDupReg, float(1.0), pregMerge);
738 Div(rstdReg, vDupReg, vMean, pregMerge);738 Div(rstdReg, vDupReg, vMean, pregMerge);
739- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr + offset, rstdReg, pregMerge);739+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr + offset, rstdReg, pregMerge);
740 }740 }
741 // unroll part741 // unroll part
742 {742 {
743 pregLoop1 = UpdateMask<float>(meanSreg1);743 pregLoop1 = UpdateMask<float>(meanSreg1);
744- DataCopy(mainA1, workAddr1 + 0);744+ LoadAlign(mainA1, workAddr1 + 0);
745- ReduceSum(vMean1, mainA1, pregLoop1);745+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
746 Muls(vMean1, vMean1, avgFactor, pregMerge1);746 Muls(vMean1, vMean1, avgFactor, pregMerge1);
747 Adds(vMean1, vMean1, epsilon, pregMerge1);747 Adds(vMean1, vMean1, epsilon, pregMerge1);
748 Sqrt(vMean1, vMean1, pregMerge1);748 Sqrt(vMean1, vMean1, pregMerge1);
749 Duplicate(vDupReg1, float(1.0), pregMerge1);749 Duplicate(vDupReg1, float(1.0), pregMerge1);
750 Div(rstdReg1, vDupReg1, vMean1, pregMerge1);750 Div(rstdReg1, vDupReg1, vMean1, pregMerge1);
751- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr1 + offset, rstdReg1, pregMerge1);751+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr1 + offset, rstdReg1, pregMerge1);
752 }752 }
753 offset += 1;753 offset += 1;
754 mainAddr += int64_t(count);754 mainAddr += int64_t(count);
@@ -767,10 +767,10 @@ __aicore__ inline void ComputeFormerImplV1MultiN(LocalTensor<T>& xLocal, LocalTe
767 }767 }
768 }768 }
769 uint32_t tailDataOffset = unrollOffset + (curRows / 2) * count;769 uint32_t tailDataOffset = unrollOffset + (curRows / 2) * count;
770- __local_mem__ T* mainAddr2 = (__ubuf__ T*)xLocal.GetPhyAddr() + tailDataOffset;770+ __ubuf__ T* mainAddr2 = (__ubuf__ T*)xLocal.GetPhyAddr() + tailDataOffset;
771- __local_mem__ T* tailAddr2 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit) + tailDataOffset;771+ __ubuf__ T* tailAddr2 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit) + tailDataOffset;
772- __local_mem__ T* masterAddr2 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile) + tailDataOffset;772+ __ubuf__ T* masterAddr2 = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile) + tailDataOffset;
773- __local_mem__ float *xFp32MainAddr2, *xFp32TailAddr2, *xFp32MasterAddr2;773+ __ubuf__ float *xFp32MainAddr2, *xFp32TailAddr2, *xFp32MasterAddr2;
774 if constexpr (is_same<T, half>::value || is_same<T, bfloat16_t>::value) {774 if constexpr (is_same<T, half>::value || is_same<T, bfloat16_t>::value) {
775 xFp32MainAddr2 = (__ubuf__ float*)xFp32.GetPhyAddr() + tailDataOffset;775 xFp32MainAddr2 = (__ubuf__ float*)xFp32.GetPhyAddr() + tailDataOffset;
776 xFp32TailAddr2 = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(powerSplit) + tailDataOffset;776 xFp32TailAddr2 = (__ubuf__ float*)xFp32.GetPhyAddr() + int64_t(powerSplit) + tailDataOffset;
@@ -796,37 +796,37 @@ __aicore__ inline void ComputeFormerImplV1MultiN(LocalTensor<T>& xLocal, LocalTe
796 Add(mainA1, mainA1, tailA1, pregLoop1);796 Add(mainA1, mainA1, tailA1, pregLoop1);
797 Add(mainB1, mainB1, tailB1, pregLoop1);797 Add(mainB1, mainB1, tailB1, pregLoop1);
798 Add(mainA1, mainA1, mainB1, pregLoop1);798 Add(mainA1, mainA1, mainB1, pregLoop1);
799- ReduceSum(vMean1, mainA1, pregLoop1);799+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
800- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);800+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);
801 }801 }
802 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {802 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
803 pregLoop1 = UpdateMask<float>(masterSreg1);803 pregLoop1 = UpdateMask<float>(masterSreg1);
804 LoadSquareMasterTile<T, true>(masterAddr2, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA1,804 LoadSquareMasterTile<T, true>(masterAddr2, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA1,
805 mainB1, pregLoop1, xFp32MasterAddr2);805 mainB1, pregLoop1, xFp32MasterAddr2);
806 Add(mainA1, mainA1, mainB1, pregLoop1);806 Add(mainA1, mainA1, mainB1, pregLoop1);
807- ReduceSum(vMean1, mainA1, pregLoop1);807+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
808- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + remainRepeats + i, vMean1, pregMerge1);808+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + remainRepeats + i, vMean1, pregMerge1);
809 }809 }
810 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();810 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
811 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {811 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
812 pregLoop1 = UpdateMask<float>(mergeSreg1);812 pregLoop1 = UpdateMask<float>(mergeSreg1);
813- DataCopy(mainA1, workAddr1 + (i * 2 + 0) * V_LENGTH);813+ LoadAlign(mainA1, workAddr1 + (i * 2 + 0) * V_LENGTH);
814- DataCopy(mainB1, workAddr1 + (i * 2 + 1) * V_LENGTH);814+ LoadAlign(mainB1, workAddr1 + (i * 2 + 1) * V_LENGTH);
815 Add(mainA1, mainA1, mainB1, pregLoop1);815 Add(mainA1, mainA1, mainB1, pregLoop1);
816- ReduceSum(vMean1, mainA1, pregLoop1);816+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
817- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);817+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr1 + i, vMean1, pregMerge1);
818 }818 }
819 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();819 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
820 {820 {
821 pregLoop1 = UpdateMask<float>(meanSreg1);821 pregLoop1 = UpdateMask<float>(meanSreg1);
822- DataCopy(mainA1, workAddr1 + 0);822+ LoadAlign(mainA1, workAddr1 + 0);
823- ReduceSum(vMean1, mainA1, pregLoop1);823+ Reduce<ReduceType::SUM>(vMean1, mainA1, pregLoop1);
824 Muls(vMean1, vMean1, avgFactor, pregMerge1);824 Muls(vMean1, vMean1, avgFactor, pregMerge1);
825 Adds(vMean1, vMean1, epsilon, pregMerge1);825 Adds(vMean1, vMean1, epsilon, pregMerge1);
826 Sqrt(vMean1, vMean1, pregMerge1);826 Sqrt(vMean1, vMean1, pregMerge1);
827 Duplicate(vDupReg1, float(1.0), pregMerge1);827 Duplicate(vDupReg1, float(1.0), pregMerge1);
828 Div(rstdReg1, vDupReg1, vMean1, pregMerge1);828 Div(rstdReg1, vDupReg1, vMean1, pregMerge1);
829- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr1 + tailOffset, rstdReg1, pregMerge1);829+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr1 + tailOffset, rstdReg1, pregMerge1);
830 }830 }
831 }831 }
832 }832 }
@@ -852,12 +852,12 @@ __aicore__ inline void ComputeFormerImplV2(LocalTensor<float>& dstLocal, LocalTe
852 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;852 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;
853 uint32_t meanSreg = meanTile;853 uint32_t meanSreg = meanTile;
854 854 
855- __local_mem__ T* mainAddr = (__ubuf__ T*)xLocal.GetPhyAddr();855+ __ubuf__ T* mainAddr = (__ubuf__ T*)xLocal.GetPhyAddr();
856- __local_mem__ T* tailAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit);856+ __ubuf__ T* tailAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(powerSplit);
857- __local_mem__ T* masterAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile);857+ __ubuf__ T* masterAddr = (__ubuf__ T*)xLocal.GetPhyAddr() + int64_t(remainTile);
858 858 
859- __local_mem__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();859+ __ubuf__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();
860- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();860+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
861 861 
862 __VEC_SCOPE__862 __VEC_SCOPE__
863 {863 {
@@ -872,32 +872,32 @@ __aicore__ inline void ComputeFormerImplV2(LocalTensor<float>& dstLocal, LocalTe
872 Add(mainA, mainA, tailA, pregLoop);872 Add(mainA, mainA, tailA, pregLoop);
873 Add(mainB, mainB, tailB, pregLoop);873 Add(mainB, mainB, tailB, pregLoop);
874 Add(mainA, mainA, mainB, pregLoop);874 Add(mainA, mainA, mainB, pregLoop);
875- ReduceSum(vMean, mainA, pregLoop);875+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
876- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);876+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);
877 }877 }
878 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {878 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
879 pregLoop = UpdateMask<float>(masterSreg);879 pregLoop = UpdateMask<float>(masterSreg);
880 LoadSquareMasterTile<T, false>(masterAddr, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA, mainB,880 LoadSquareMasterTile<T, false>(masterAddr, (i * 2 + 0) * V_LENGTH, (i * 2 + 1) * V_LENGTH, mainA, mainB,
881 pregLoop);881 pregLoop);
882 Add(mainA, mainA, mainB, pregLoop);882 Add(mainA, mainA, mainB, pregLoop);
883- ReduceSum(vMean, mainA, pregLoop);883+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
884- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + remainRepeats + i, vMean, pregMerge);884+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + remainRepeats + i, vMean, pregMerge);
885 }885 }
886 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();886 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
887 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {887 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
888 pregLoop = UpdateMask<float>(mergeSreg);888 pregLoop = UpdateMask<float>(mergeSreg);
889- DataCopy(mainA, workAddr + (i * 2 + 0) * V_LENGTH);889+ LoadAlign(mainA, workAddr + (i * 2 + 0) * V_LENGTH);
890- DataCopy(mainB, workAddr + (i * 2 + 1) * V_LENGTH);890+ LoadAlign(mainB, workAddr + (i * 2 + 1) * V_LENGTH);
891 Add(mainA, mainA, mainB, pregLoop);891 Add(mainA, mainA, mainB, pregLoop);
892- ReduceSum(vMean, mainA, pregLoop);892+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
893- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);893+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + i, vMean, pregMerge);
894 }894 }
895 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();895 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
896 {896 {
897 pregLoop = UpdateMask<float>(meanSreg);897 pregLoop = UpdateMask<float>(meanSreg);
898- DataCopy(mainA, workAddr + 0);898+ LoadAlign(mainA, workAddr + 0);
899- ReduceSum(vMean, mainA, pregLoop);899+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
900- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMean, pregMerge);900+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMean, pregMerge);
901 }901 }
902 }902 }
903}903}
Mnorm/rms_norm/op_kernel/arch35/rms_norm_regbase_perf.h+5-5
@@ -159,10 +159,10 @@ private:
159 __aicore__ inline void ComputeY(LocalTensor<DX> xLocal, LocalTensor<DG> gammaLocal, LocalTensor<float> rstdLocal,159 __aicore__ inline void ComputeY(LocalTensor<DX> xLocal, LocalTensor<DG> gammaLocal, LocalTensor<float> rstdLocal,
160 LocalTensor<DX> yLocal, uint64_t curUbFactor)160 LocalTensor<DX> yLocal, uint64_t curUbFactor)
161 {161 {
162- __local_mem__ DX* xLocalAddr = (__local_mem__ DX*)xLocal.GetPhyAddr();162+ __ubuf__ DX* xLocalAddr = (__ubuf__ DX*)xLocal.GetPhyAddr();
163- __local_mem__ DG* gammaLocalUbAddr = (__local_mem__ DG*)gammaLocal.GetPhyAddr();163+ __ubuf__ DG* gammaLocalUbAddr = (__ubuf__ DG*)gammaLocal.GetPhyAddr();
164- __local_mem__ float* rstdLocalUbAddr = (__local_mem__ float*)rstdLocal.GetPhyAddr();164+ __ubuf__ float* rstdLocalUbAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
165- __local_mem__ DX* yLocalUbAddr = (__local_mem__ DX*)yLocal.GetPhyAddr();165+ __ubuf__ DX* yLocalUbAddr = (__ubuf__ DX*)yLocal.GetPhyAddr();
166 166 
167 uint32_t colNum = static_cast<uint32_t>(numCol);167 uint32_t colNum = static_cast<uint32_t>(numCol);
168 uint16_t curAloops = static_cast<uint16_t>(curUbFactor);168 uint16_t curAloops = static_cast<uint16_t>(curUbFactor);
@@ -179,7 +179,7 @@ private:
179 179 
180 for (uint16_t i = 0; i < curAloops; i++) {180 for (uint16_t i = 0; i < curAloops; i++) {
181 uint32_t sregElewiseNum = numCol;181 uint32_t sregElewiseNum = numCol;
182- DataCopy<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdLocalUbAddr + i);182+ LoadAlign<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdLocalUbAddr + i);
183 for (uint16_t j = 0; j < colLoops; j++) {183 for (uint16_t j = 0; j < colLoops; j++) {
184 MaskReg pregCurLoop = UpdateMask<float>(sregElewiseNum);184 MaskReg pregCurLoop = UpdateMask<float>(sregElewiseNum);
185 LoadRegForDtype(xLocalAddr, xReg, pregCurLoop, (i * colNumAlign + j * VectorLenB32));185 LoadRegForDtype(xLocalAddr, xReg, pregCurLoop, (i * colNumAlign + j * VectorLenB32));
Mnorm/rms_norm_dynamic_mx_quant/op_kernel/arch35/rms_norm_dynamic_mx_quant_common.h+75-75
@@ -81,26 +81,26 @@ constexpr uint32_t FP8_E4M3_MAX = 0x3b124925; // 1/448的float32表示 448是E4M
81constexpr uint16_t INVALID_FLOAT16 = 0x7c00;81constexpr uint16_t INVALID_FLOAT16 = 0x7c00;
82 82 
83template <typename T>83template <typename T>
84-__aicore__ inline void LoadTensorForDtypeT(__local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)84+__aicore__ inline void LoadTensorForDtypeT(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)
85{85{
86 if constexpr (IsSameType<T, float>::value) {86 if constexpr (IsSameType<T, float>::value) {
87- DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset);87+ LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset);
88 } else { // fp16、bf1688 } else { // fp16、bf16
89 RegTensor<T> xFp16;89 RegTensor<T> xFp16;
90- DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset));90+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset));
91 Cast<float, T, castTraitB162B32>(dst, xFp16, preg);91 Cast<float, T, castTraitB162B32>(dst, xFp16, preg);
92 }92 }
93}93}
94 94 
95template <typename T>95template <typename T>
96-__aicore__ inline void StoreTensorForDtypeT(__local_mem__ T* dst, RegTensor<float>& src, MaskReg& preg, uint32_t offset)96+__aicore__ inline void StoreTensorForDtypeT(__ubuf__ T* dst, RegTensor<float>& src, MaskReg& preg, uint32_t offset)
97{97{
98 if constexpr (IsSameType<T, float>::value) {98 if constexpr (IsSameType<T, float>::value) {
99- DataCopy<T, StoreDist::DIST_NORM>(dst + offset, src, preg);99+ StoreAlign<T, StoreDist::DIST_NORM>(dst + offset, src, preg);
100 } else {100 } else {
101 RegTensor<T> xOut;101 RegTensor<T> xOut;
102 Cast<T, float, castTraitB322B16>(xOut, src, preg);102 Cast<T, float, castTraitB322B16>(xOut, src, preg);
103- DataCopy<T, StoreDist::DIST_PACK_B32>(dst + offset, xOut, preg);103+ StoreAlign<T, StoreDist::DIST_PACK_B32>(dst + offset, xOut, preg);
104 }104 }
105}105}
106 106 
@@ -137,12 +137,12 @@ __aicore__ inline void ComputeData(__ubuf__ T1* srcAddr, __ubuf__ uint16_t* half
137 dataMask4 = AscendC::MicroAPI::CreateMask<T1>();137 dataMask4 = AscendC::MicroAPI::CreateMask<T1>();
138 dataMask5 = AscendC::MicroAPI::CreateMask<T2>();138 dataMask5 = AscendC::MicroAPI::CreateMask<T2>();
139 for (uint16_t i = 0; i < loopNum; i++) {139 for (uint16_t i = 0; i < loopNum; i++) {
140- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,140+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
141- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,141+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,
142- VL_B16 * NUM_TWO);142+ VL_B16 * NUM_TWO);
143- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,143+ AscendC::MicroAPI::LoadAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
144- AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr,144+ AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr,
145- elementAfterReduce);145+ elementAfterReduce);
146 if constexpr (IsSameType<T1, half>::value) {146 if constexpr (IsSameType<T1, half>::value) {
147 AscendC::MicroAPI::Cast<float, T1, castTraitZero>(vdExp0FP32Zero, vdExp0, dataMask1);147 AscendC::MicroAPI::Cast<float, T1, castTraitZero>(vdExp0FP32Zero, vdExp0, dataMask1);
148 AscendC::MicroAPI::Cast<float, T1, castTraitOne>(vdExp0FP32One, vdExp0, dataMask1);148 AscendC::MicroAPI::Cast<float, T1, castTraitOne>(vdExp0FP32One, vdExp0, dataMask1);
@@ -180,8 +180,8 @@ __aicore__ inline void ComputeData(__ubuf__ T1* srcAddr, __ubuf__ uint16_t* half
180 (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,180 (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp0FP8Zero,
181 (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp1FP8One, dataMask5);181 (AscendC::MicroAPI::RegTensor<uint8_t>&)vdExp1FP8One, dataMask5);
182 182 
183- AscendC::MicroAPI::DataCopy<int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,183+ AscendC::MicroAPI::StoreAlign<int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
184- AscendC::MicroAPI::StoreDist::DIST_NORM_B8>(184+ AscendC::MicroAPI::StoreDist::DIST_NORM_B8>(
185 outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP8Zero, OUT_ALL, dataMask5);185 outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP8Zero, OUT_ALL, dataMask5);
186 }186 }
187 }187 }
@@ -215,13 +215,13 @@ __aicore__ inline void ComputeMaxExpOCP(__ubuf__ T* srcAddr, __ubuf__ uint16_t*
215 AscendC::MicroAPI::MaskReg scaleMask2;215 AscendC::MicroAPI::MaskReg scaleMask2;
216 AscendC::MicroAPI::MaskReg invalidDataMask0;216 AscendC::MicroAPI::MaskReg invalidDataMask0;
217 AscendC::MicroAPI::MaskReg invalidDataMask1;217 AscendC::MicroAPI::MaskReg invalidDataMask1;
218- AscendC::MicroAPI::UnalignReg u1;218+ AscendC::MicroAPI::UnalignRegForStore u1;
219 for (uint16_t i = 0; i < loopNum; i++) {219 for (uint16_t i = 0; i < loopNum; i++) {
220 scaleMask1 = AscendC::MicroAPI::UpdateMask<T>(totalCountInUB);220 scaleMask1 = AscendC::MicroAPI::UpdateMask<T>(totalCountInUB);
221 scaleMask2 = AscendC::MicroAPI::UpdateMask<T>(totalCountInUB);221 scaleMask2 = AscendC::MicroAPI::UpdateMask<T>(totalCountInUB);
222- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,222+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
223- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,223+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,
224- VL_B16 * NUM_TWO);224+ VL_B16 * NUM_TWO);
225 if constexpr (IsSameType<T, half>::value) {225 if constexpr (IsSameType<T, half>::value) {
226 AscendC::MicroAPI::And(vdExpSelect0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, invalidMaskFP16,226 AscendC::MicroAPI::And(vdExpSelect0, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, invalidMaskFP16,
227 scaleMask1);227 scaleMask1);
@@ -247,12 +247,12 @@ __aicore__ inline void ComputeMaxExpOCP(__ubuf__ T* srcAddr, __ubuf__ uint16_t*
247 }247 }
248 248 
249 AscendC::MicroAPI::Max(vdMaxExp, vdExpExtract0, vdExpExtract1, scaleMask1);249 AscendC::MicroAPI::Max(vdMaxExp, vdExpExtract0, vdExpExtract1, scaleMask1);
250- AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, scaleMask1);250+ AscendC::MicroAPI::ReduceDataBlock<ReduceType::MAX>(vdMaxExp, vdMaxExp, scaleMask1);
251 251 
252- AscendC::MicroAPI::DataCopyUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(252+ AscendC::MicroAPI::StoreUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
253 maxExpAddr, vdMaxExp, u1, elementAfterReduce);253 maxExpAddr, vdMaxExp, u1, elementAfterReduce);
254 }254 }
255- AscendC::MicroAPI::DataCopyUnAlignPost(maxExpAddr, u1, 0);255+ AscendC::MicroAPI::StoreUnAlignPost(maxExpAddr, u1, 0);
256 }256 }
257 return;257 return;
258}258}
@@ -300,8 +300,8 @@ __aicore__ inline void ComputeScaleOCP(__ubuf__ uint16_t* maxExpAddr, __ubuf__ u
300 AscendC::MicroAPI::Duplicate(specialExpRegTensor, SPECIAL_EXP_THRESHOLD);300 AscendC::MicroAPI::Duplicate(specialExpRegTensor, SPECIAL_EXP_THRESHOLD);
301 for (uint16_t i = 0; i < loopNumScale; i++) {301 for (uint16_t i = 0; i < loopNumScale; i++) {
302 preMaskScale = AscendC::MicroAPI::UpdateMask<uint16_t>(totalScaleInUB);302 preMaskScale = AscendC::MicroAPI::UpdateMask<uint16_t>(totalScaleInUB);
303- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(vdMaxExp,303+ AscendC::MicroAPI::LoadAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
304- maxExpAddr, VL_B16);304+ vdMaxExp, maxExpAddr, VL_B16);
305 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(cmpResult, vdMaxExp, expMask, preMaskScale); // INF/NAN305 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(cmpResult, vdMaxExp, expMask, preMaskScale); // INF/NAN
306 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::LE>(invalidDataMask, vdMaxExp, maxExpValue, preMaskScale);306 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::LE>(invalidDataMask, vdMaxExp, maxExpValue, preMaskScale);
307 307 
@@ -312,9 +312,9 @@ __aicore__ inline void ComputeScaleOCP(__ubuf__ uint16_t* maxExpAddr, __ubuf__ u
312 312 
313 AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, fp8NanRegTensor, cmpResult);313 AscendC::MicroAPI::Select<uint16_t>(scaleValue, scaleValue, fp8NanRegTensor, cmpResult);
314 314 
315- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,315+ AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
316- AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(mxScaleLocalAddr, scaleValue,316+ AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(mxScaleLocalAddr, scaleValue,
317- VL_B16 / NUM_TWO, preMaskScale);317+ VL_B16 / NUM_TWO, preMaskScale);
318 318 
319 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(zeroMask, sharedExp, zeroRegTensor, preMaskScale);319 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::NE>(zeroMask, sharedExp, zeroRegTensor, preMaskScale);
320 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::EQ>(specialDataMask, sharedExp, scaleBias, preMaskScale);320 AscendC::MicroAPI::Compare<uint16_t, CMPMODE::EQ>(specialDataMask, sharedExp, scaleBias, preMaskScale);
@@ -323,7 +323,7 @@ __aicore__ inline void ComputeScaleOCP(__ubuf__ uint16_t* maxExpAddr, __ubuf__ u
323 AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, zeroRegTensor, zeroMask);323 AscendC::MicroAPI::Select<uint16_t>(halfScale, halfScale, zeroRegTensor, zeroMask);
324 AscendC::MicroAPI::Select<uint16_t>(halfScale, specialExpRegTensor, halfScale, specialDataMask);324 AscendC::MicroAPI::Select<uint16_t>(halfScale, specialExpRegTensor, halfScale, specialDataMask);
325 325 
326- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(326+ AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
327 halfScaleLocalAddr, halfScale, VL_B16, preMaskScale);327 halfScaleLocalAddr, halfScale, VL_B16, preMaskScale);
328 }328 }
329 }329 }
@@ -344,23 +344,23 @@ __aicore__ inline void ComputeMaxExpcuBLAS(__ubuf__ T* srcAddr, __ubuf__ uint16_
344 AscendC::MicroAPI::Duplicate(absMask16Bit, ABS_MASK_FOR_16BIT);344 AscendC::MicroAPI::Duplicate(absMask16Bit, ABS_MASK_FOR_16BIT);
345 AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;345 AscendC::MicroAPI::RegTensor<uint16_t> vdMaxExp;
346 AscendC::MicroAPI::MaskReg scaleMask1;346 AscendC::MicroAPI::MaskReg scaleMask1;
347- AscendC::MicroAPI::UnalignReg u1;347+ AscendC::MicroAPI::UnalignRegForStore u1;
348 for (uint16_t i = 0; i < loopNum; i++) {348 for (uint16_t i = 0; i < loopNum; i++) {
349 scaleMask1 = AscendC::MicroAPI::UpdateMask<T>(totalCountInUB);349 scaleMask1 = AscendC::MicroAPI::UpdateMask<T>(totalCountInUB);
350- AscendC::MicroAPI::DataCopy<T, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,350+ AscendC::MicroAPI::LoadAlign<T, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
351- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,351+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,
352- VL_B16 * NUM_TWO);352+ VL_B16 * NUM_TWO);
353 AscendC::MicroAPI::And((AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,353 AscendC::MicroAPI::And((AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,
354 (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, absMask16Bit, scaleMask1);354 (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0, absMask16Bit, scaleMask1);
355 AscendC::MicroAPI::And((AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1,355 AscendC::MicroAPI::And((AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1,
356 (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, absMask16Bit, scaleMask1);356 (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, absMask16Bit, scaleMask1);
357 AscendC::MicroAPI::Max(vdMaxExp, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,357 AscendC::MicroAPI::Max(vdMaxExp, (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp0,
358 (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, scaleMask1);358 (AscendC::MicroAPI::RegTensor<uint16_t>&)vdExp1, scaleMask1);
359- AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, scaleMask1);359+ AscendC::MicroAPI::ReduceDataBlock<ReduceType::MAX>(vdMaxExp, vdMaxExp, scaleMask1);
360- AscendC::MicroAPI::DataCopyUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(360+ AscendC::MicroAPI::StoreUnAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
361 maxExpAddr, vdMaxExp, u1, elementAfterReduce);361 maxExpAddr, vdMaxExp, u1, elementAfterReduce);
362 }362 }
363- AscendC::MicroAPI::DataCopyUnAlignPost(maxExpAddr, u1, 0);363+ AscendC::MicroAPI::StoreUnAlignPost(maxExpAddr, u1, 0);
364 }364 }
365 return;365 return;
366}366}
@@ -416,8 +416,8 @@ __aicore__ inline void ComputeScalecuBLAS(__ubuf__ uint16_t* maxExpAddr, __ubuf_
416 preMaskScale = AscendC::MicroAPI::CreateMask<uint32_t>();416 preMaskScale = AscendC::MicroAPI::CreateMask<uint32_t>();
417 maskHalf = AscendC::MicroAPI::CreateMask<uint16_t>();417 maskHalf = AscendC::MicroAPI::CreateMask<uint16_t>();
418 for (uint16_t i = 0; i < loopNumScale4NV; i++) {418 for (uint16_t i = 0; i < loopNumScale4NV; i++) {
419- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,419+ AscendC::MicroAPI::LoadAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
420- AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(max16, maxExpAddr, VL_FP32);420+ AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(max16, maxExpAddr, VL_FP32);
421 421 
422 AscendC::MicroAPI::Cast<float, T1, castTraitHalf2Float>(422 AscendC::MicroAPI::Cast<float, T1, castTraitHalf2Float>(
423 (AscendC::MicroAPI::RegTensor<float>&)max32, (AscendC::MicroAPI::RegTensor<T1>&)max16, preMaskScale);423 (AscendC::MicroAPI::RegTensor<float>&)max32, (AscendC::MicroAPI::RegTensor<T1>&)max16, preMaskScale);
@@ -430,16 +430,16 @@ __aicore__ inline void ComputeScalecuBLAS(__ubuf__ uint16_t* maxExpAddr, __ubuf_
430 AscendC::MicroAPI::ShiftRights(exp32, max32, SHR_NUM_FOR_FP32, preMaskScale);430 AscendC::MicroAPI::ShiftRights(exp32, max32, SHR_NUM_FOR_FP32, preMaskScale);
431 AscendC::MicroAPI::And(man32, max32, manMaskFP32, preMaskScale);431 AscendC::MicroAPI::And(man32, max32, manMaskFP32, preMaskScale);
432 432 
433- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p0, exp32, zeroForAll, preMaskScale);433+ AscendC::MicroAPI::Compares<uint32_t, CMPMODE::GT>(p0, exp32, zeroForAll, preMaskScale);
434- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::LT>(p1, exp32, Exp254, preMaskScale);434+ AscendC::MicroAPI::Compares<uint32_t, CMPMODE::LT>(p1, exp32, Exp254, preMaskScale);
435- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p2, man32, zeroForAll, preMaskScale);435+ AscendC::MicroAPI::Compares<uint32_t, CMPMODE::GT>(p2, man32, zeroForAll, preMaskScale);
436- AscendC::MicroAPI::MaskAnd(p0, p0, p1, preMaskScale);436+ AscendC::MicroAPI::And(p0, p0, p1, preMaskScale);
437- AscendC::MicroAPI::MaskAnd(p0, p0, p2, preMaskScale);437+ AscendC::MicroAPI::And(p0, p0, p2, preMaskScale);
438 438 
439- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::EQ>(p1, exp32, zeroForAll, preMaskScale);439+ AscendC::MicroAPI::Compares<uint32_t, CMPMODE::EQ>(p1, exp32, zeroForAll, preMaskScale);
440- AscendC::MicroAPI::CompareScalar<uint32_t, CMPMODE::GT>(p2, man32, halfForMan, preMaskScale);440+ AscendC::MicroAPI::Compares<uint32_t, CMPMODE::GT>(p2, man32, halfForMan, preMaskScale);
441- AscendC::MicroAPI::MaskAnd(p1, p1, p2, preMaskScale);441+ AscendC::MicroAPI::And(p1, p1, p2, preMaskScale);
442- AscendC::MicroAPI::MaskOr(p0, p0, p1, preMaskScale);442+ AscendC::MicroAPI::Or(p0, p0, p1, preMaskScale);
443 443 
444 AscendC::MicroAPI::Adds(expAddOne32, exp32, 1, preMaskScale);444 AscendC::MicroAPI::Adds(expAddOne32, exp32, 1, preMaskScale);
445 AscendC::MicroAPI::Select(extractExp, expAddOne32, exp32, p0);445 AscendC::MicroAPI::Select(extractExp, expAddOne32, exp32, p0);
@@ -447,9 +447,9 @@ __aicore__ inline void ComputeScalecuBLAS(__ubuf__ uint16_t* maxExpAddr, __ubuf_
447 AscendC::MicroAPI::Select<uint32_t>(extractExp, extractExp, zeroRegTensor32, zeroMask);447 AscendC::MicroAPI::Select<uint32_t>(extractExp, extractExp, zeroRegTensor32, zeroMask);
448 AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(expOut, extractExp);448 AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(expOut, extractExp);
449 449 
450- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,450+ AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
451- AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(mxScaleLocalAddr, expOut,451+ AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(mxScaleLocalAddr, expOut,
452- VL_FP32 / NUM_TWO, maskHalf);452+ VL_FP32 / NUM_TWO, maskHalf);
453 453 
454 AscendC::MicroAPI::ShiftLefts(extractExp, extractExp, SHR_NUM_FOR_BF16, preMaskScale);454 AscendC::MicroAPI::ShiftLefts(extractExp, extractExp, SHR_NUM_FOR_BF16, preMaskScale);
455 AscendC::MicroAPI::Sub(halfScale, scaleBias, extractExp, preMaskScale);455 AscendC::MicroAPI::Sub(halfScale, scaleBias, extractExp, preMaskScale);
@@ -457,7 +457,7 @@ __aicore__ inline void ComputeScalecuBLAS(__ubuf__ uint16_t* maxExpAddr, __ubuf_
457 AscendC::MicroAPI::Select<uint32_t>(halfScale, halfScale, zeroRegTensor32, zeroMask);457 AscendC::MicroAPI::Select<uint32_t>(halfScale, halfScale, zeroRegTensor32, zeroMask);
458 AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(recExpOut, halfScale);458 AscendC::MicroAPI::Pack<uint16_t, uint32_t, AscendC::MicroAPI::HighLowPart::LOWEST>(recExpOut, halfScale);
459 459 
460- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(460+ AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE>(
461 halfScaleLocalAddr, recExpOut, VL_FP32, maskHalf);461 halfScaleLocalAddr, recExpOut, VL_FP32, maskHalf);
462 }462 }
463 }463 }
@@ -481,7 +481,7 @@ __aicore__ inline void ComputeFP4FromHalf(MicroAPI::RegTensor<float>& Reg)
481 MicroAPI::Compare<int32_t, CMPMODE::EQ>(negInfMask, (MicroAPI::RegTensor<int32_t>&)Reg, negZero, pregAll32);481 MicroAPI::Compare<int32_t, CMPMODE::EQ>(negInfMask, (MicroAPI::RegTensor<int32_t>&)Reg, negZero, pregAll32);
482 if constexpr (IsSameType<T2, fp4x2_e1m2_t>::value) {482 if constexpr (IsSameType<T2, fp4x2_e1m2_t>::value) {
483 MicroAPI::Muls(Reg, Reg, FOUR, pregAll32);483 MicroAPI::Muls(Reg, Reg, FOUR, pregAll32);
484- MicroAPI::CompareScalar<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);484+ MicroAPI::Compares<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);
485 MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);485 MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);
486 MicroAPI::Muls(Reg, Reg, ONE_FOURTH, pregAll32);486 MicroAPI::Muls(Reg, Reg, ONE_FOURTH, pregAll32);
487 } else {487 } else {
@@ -498,13 +498,13 @@ __aicore__ inline void ComputeFP4FromHalf(MicroAPI::RegTensor<float>& Reg)
498 MicroAPI::Mul(Reg, Reg, (MicroAPI::RegTensor<float>&)exp1FP32, pregAll32);498 MicroAPI::Mul(Reg, Reg, (MicroAPI::RegTensor<float>&)exp1FP32, pregAll32);
499 MicroAPI::Adds(exp0FP32, exp0FP32, FP32_BIAS, pregAll32);499 MicroAPI::Adds(exp0FP32, exp0FP32, FP32_BIAS, pregAll32);
500 MicroAPI::ShiftLefts(exp0FP32, exp0FP32, SHR_NUM_FOR_FP32, pregAll32);500 MicroAPI::ShiftLefts(exp0FP32, exp0FP32, SHR_NUM_FOR_FP32, pregAll32);
501- MicroAPI::CompareScalar<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);501+ MicroAPI::Compares<float, CMPMODE::LT>(specialMask, Reg, 0, pregAll32);
502 MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);502 MicroAPI::Truncate<float, roundMode>(Reg, Reg, pregAll32);
503 MicroAPI::Mul(Reg, Reg, (MicroAPI::RegTensor<float>&)exp0FP32, pregAll32);503 MicroAPI::Mul(Reg, Reg, (MicroAPI::RegTensor<float>&)exp0FP32, pregAll32);
504 }504 }
505- MicroAPI::CompareScalar<float, CMPMODE::EQ>(zeroMask, Reg, 0, pregAll32);505+ MicroAPI::Compares<float, CMPMODE::EQ>(zeroMask, Reg, 0, pregAll32);
506- MicroAPI::MaskAnd(zeroMask, specialMask, zeroMask, pregAll32);506+ MicroAPI::And(zeroMask, specialMask, zeroMask, pregAll32);
507- MicroAPI::MaskOr(zeroMask, negInfMask, zeroMask, pregAll32);507+ MicroAPI::Or(zeroMask, negInfMask, zeroMask, pregAll32);
508 MicroAPI::Select<int32_t>((MicroAPI::RegTensor<int32_t>&)Reg, negZero, (MicroAPI::RegTensor<int32_t>&)Reg,508 MicroAPI::Select<int32_t>((MicroAPI::RegTensor<int32_t>&)Reg, negZero, (MicroAPI::RegTensor<int32_t>&)Reg,
509 zeroMask);509 zeroMask);
510}510}
@@ -528,9 +528,9 @@ __aicore__ inline void FP16Convert(AscendC::MicroAPI::RegTensor<half>& output,
528 AscendC::MicroAPI::Duplicate(specialValueTensor, specialValue);528 AscendC::MicroAPI::Duplicate(specialValueTensor, specialValue);
529 AscendC::MicroAPI::Duplicate(newMantissa, NEW_MANTISSA);529 AscendC::MicroAPI::Duplicate(newMantissa, NEW_MANTISSA);
530 AscendC::MicroAPI::And(andResult, (AscendC::MicroAPI::RegTensor<uint16_t>&)input, specialValueTensor, mask);530 AscendC::MicroAPI::And(andResult, (AscendC::MicroAPI::RegTensor<uint16_t>&)input, specialValueTensor, mask);
531- AscendC::MicroAPI::CompareScalar<uint16_t, CMPMODE::GT>(nonzeroMask, andResult, 0, mask);531+ AscendC::MicroAPI::Compares<uint16_t, CMPMODE::GT>(nonzeroMask, andResult, 0, mask);
532- AscendC::MicroAPI::CompareScalar<uint16_t, CMPMODE::LT>(specialMask, andResult, NEW_MANTISSA, mask);532+ AscendC::MicroAPI::Compares<uint16_t, CMPMODE::LT>(specialMask, andResult, NEW_MANTISSA, mask);
533- AscendC::MicroAPI::MaskAnd(specialMask, specialMask, nonzeroMask, mask);533+ AscendC::MicroAPI::And(specialMask, specialMask, nonzeroMask, mask);
534 AscendC::MicroAPI::Or(newValue, (AscendC::MicroAPI::RegTensor<uint16_t>&)input, newMantissa, mask);534 AscendC::MicroAPI::Or(newValue, (AscendC::MicroAPI::RegTensor<uint16_t>&)input, newMantissa, mask);
535 AscendC::MicroAPI::Select<uint16_t>((AscendC::MicroAPI::RegTensor<uint16_t>&)output, newValue,535 AscendC::MicroAPI::Select<uint16_t>((AscendC::MicroAPI::RegTensor<uint16_t>&)output, newValue,
536 (AscendC::MicroAPI::RegTensor<uint16_t>&)input, specialMask);536 (AscendC::MicroAPI::RegTensor<uint16_t>&)input, specialMask);
@@ -557,12 +557,12 @@ __aicore__ inline void ComputeDataMxfp4General(__ubuf__ T1* srcAddr, __ubuf__ ui
557 557 
558 for (uint16_t i = 0; i < loopNum; i++) {558 for (uint16_t i = 0; i < loopNum; i++) {
559 dataMask1 = AscendC::MicroAPI::UpdateMask<T1>(totalCountInUB);559 dataMask1 = AscendC::MicroAPI::UpdateMask<T1>(totalCountInUB);
560- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,560+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
561- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,561+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,
562- VL_B16 * NUM_TWO);562+ VL_B16 * NUM_TWO);
563- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,563+ AscendC::MicroAPI::LoadAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
564- AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr,564+ AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr,
565- VL_BLOCK_NUM);565+ VL_BLOCK_NUM);
566 566 
567 if constexpr (IsSameType<T1, half>::value) {567 if constexpr (IsSameType<T1, half>::value) {
568 if constexpr (roundMode == RoundMode::CAST_RINT) {568 if constexpr (roundMode == RoundMode::CAST_RINT) {
@@ -588,9 +588,9 @@ __aicore__ inline void ComputeDataMxfp4General(__ubuf__ T1* srcAddr, __ubuf__ ui
588 AscendC::MicroAPI::Cast<T2, T1, castTraitRM<roundMode>>(vdExp1FP4, vdExp1, dataMask1);588 AscendC::MicroAPI::Cast<T2, T1, castTraitRM<roundMode>>(vdExp1FP4, vdExp1, dataMask1);
589 }589 }
590 590 
591- MicroAPI::DataCopy<int8_t, MicroAPI::PostLiteral::POST_MODE_UPDATE, MicroAPI::StoreDist::DIST_PACK4_B32>(591+ MicroAPI::StoreAlign<int8_t, MicroAPI::PostLiteral::POST_MODE_UPDATE, MicroAPI::StoreDist::DIST_PACK4_B32>(
592 outLocalAddr, (MicroAPI::RegTensor<int8_t>&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);592 outLocalAddr, (MicroAPI::RegTensor<int8_t>&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
593- MicroAPI::DataCopy<int8_t, MicroAPI::PostLiteral::POST_MODE_UPDATE, MicroAPI::StoreDist::DIST_PACK4_B32>(593+ MicroAPI::StoreAlign<int8_t, MicroAPI::PostLiteral::POST_MODE_UPDATE, MicroAPI::StoreDist::DIST_PACK4_B32>(
594 outLocalAddr, (MicroAPI::RegTensor<int8_t>&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);594 outLocalAddr, (MicroAPI::RegTensor<int8_t>&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
595 }595 }
596 }596 }
@@ -626,12 +626,12 @@ __aicore__ inline void ComputeDataMxfp4Optimize(__ubuf__ T1* srcAddr, __ubuf__ u
626 626 
627 for (uint16_t i = 0; i < loopNum; i++) {627 for (uint16_t i = 0; i < loopNum; i++) {
628 dataMask1 = AscendC::MicroAPI::UpdateMask<T1>(totalCountInUB);628 dataMask1 = AscendC::MicroAPI::UpdateMask<T1>(totalCountInUB);
629- AscendC::MicroAPI::DataCopy<T1, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,629+ AscendC::MicroAPI::LoadAlign<T1, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
630- AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,630+ AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr,
631- VL_B16 * NUM_TWO);631+ VL_B16 * NUM_TWO);
632- AscendC::MicroAPI::DataCopy<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,632+ AscendC::MicroAPI::LoadAlign<uint16_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
633- AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr,633+ AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr,
634- VL_BLOCK_NUM);634+ VL_BLOCK_NUM);
635 635 
636 if constexpr (IsSameType<T1, half>::value) {636 if constexpr (IsSameType<T1, half>::value) {
637 MicroAPI::Cast<float, bfloat16_t, castTraitF16toFp32Zero>(637 MicroAPI::Cast<float, bfloat16_t, castTraitF16toFp32Zero>(
@@ -684,11 +684,11 @@ __aicore__ inline void ComputeDataMxfp4Optimize(__ubuf__ T1* srcAddr, __ubuf__ u
684 AscendC::MicroAPI::Cast<T2, T1, castTraitRM<roundMode>>(vdExp1FP4, vdExp1, dataMask1);684 AscendC::MicroAPI::Cast<T2, T1, castTraitRM<roundMode>>(vdExp1FP4, vdExp1, dataMask1);
685 }685 }
686 686 
687- AscendC::MicroAPI::DataCopy<int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,687+ AscendC::MicroAPI::StoreAlign<int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
688- AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(688+ AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(
689 outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);689 outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
690- AscendC::MicroAPI::DataCopy<int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,690+ AscendC::MicroAPI::StoreAlign<int8_t, AscendC::MicroAPI::PostLiteral::POST_MODE_UPDATE,
691- AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(691+ AscendC::MicroAPI::StoreDist::DIST_PACK4_B32>(
692 outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);692 outLocalAddr, (AscendC::MicroAPI::RegTensor<int8_t>&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask1);
693 }693 }
694 }694 }
Mnorm/rms_norm_dynamic_mx_quant/op_kernel/arch35/rms_norm_dynamic_mx_quant_full_load.h+8-8
@@ -224,15 +224,15 @@ public:
224 LocalTensor<T_GAMMA> betaLocal, LocalTensor<float> rstdLocal,224 LocalTensor<T_GAMMA> betaLocal, LocalTensor<float> rstdLocal,
225 LocalTensor<T_X> yLocal, int64_t curM)225 LocalTensor<T_X> yLocal, int64_t curM)
226 {226 {
227- __local_mem__ T_X* xLocalAddr = (__local_mem__ T_X*)xLocal.GetPhyAddr();227+ __ubuf__ T_X* xLocalAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
228- __local_mem__ T_GAMMA* gammaLocalUbAddr = (__local_mem__ T_GAMMA*)gammaLocal.GetPhyAddr();228+ __ubuf__ T_GAMMA* gammaLocalUbAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
229- __local_mem__ T_GAMMA* betaLocalUbAddr;229+ __ubuf__ T_GAMMA* betaLocalUbAddr;
230 if constexpr (hasInputBeta) {230 if constexpr (hasInputBeta) {
231- betaLocalUbAddr = (__local_mem__ T_GAMMA*)betaLocal.GetPhyAddr();231+ betaLocalUbAddr = (__ubuf__ T_GAMMA*)betaLocal.GetPhyAddr();
232 }232 }
233 233 
234- __local_mem__ float* rstdLocalUbAddr = (__local_mem__ float*)rstdLocal.GetPhyAddr();234+ __ubuf__ float* rstdLocalUbAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
235- __local_mem__ T_X* yLocalUbAddr = (__local_mem__ T_X*)yLocal.GetPhyAddr();235+ __ubuf__ T_X* yLocalUbAddr = (__ubuf__ T_X*)yLocal.GetPhyAddr();
236 uint32_t nNum = static_cast<uint32_t>(tilingData_->numN);236 uint32_t nNum = static_cast<uint32_t>(tilingData_->numN);
237 uint16_t mloops = static_cast<uint16_t>(curM);237 uint16_t mloops = static_cast<uint16_t>(curM);
238 uint16_t nloops = static_cast<uint16_t>(ops::CeilDiv(nNum, VL_FP32));238 uint16_t nloops = static_cast<uint16_t>(ops::CeilDiv(nNum, VL_FP32));
@@ -251,7 +251,7 @@ public:
251 AscendC::MicroAPI::MaskReg251 AscendC::MicroAPI::MaskReg
252 pregFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();252 pregFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
253 for (uint16_t i = 0; i < mloops; i++) {253 for (uint16_t i = 0; i < mloops; i++) {
254- DataCopy<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdLocalUbAddr + i);254+ LoadAlign<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdLocalUbAddr + i);
255 uint32_t xElemOffset = i * xInputStride;255 uint32_t xElemOffset = i * xInputStride;
256 LoadTensorForDtypeT<T_X>(xLocalAddr, xReg, pregMask, xElemOffset);256 LoadTensorForDtypeT<T_X>(xLocalAddr, xReg, pregMask, xElemOffset);
257 Mul(yReg, xReg, RstdReg, pregMask);257 Mul(yReg, xReg, RstdReg, pregMask);
@@ -280,7 +280,7 @@ public:
280 pregFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();280 pregFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
281 for (uint16_t i = 0; i < mloops; i++) {281 for (uint16_t i = 0; i < mloops; i++) {
282 uint32_t sreg = nNum;282 uint32_t sreg = nNum;
283- DataCopy<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdLocalUbAddr + i);283+ LoadAlign<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdLocalUbAddr + i);
284 for (uint16_t j = 0; j < nloops; j++) {284 for (uint16_t j = 0; j < nloops; j++) {
285 pregMask = UpdateMask<float>(sreg);285 pregMask = UpdateMask<float>(sreg);
286 uint32_t gammaElemOffset = j * VL_FP32;286 uint32_t gammaElemOffset = j * VL_FP32;
Mnorm/rms_norm_dynamic_mx_quant/op_kernel/arch35/rms_norm_dynamic_mx_quant_recompute.h+26-26
@@ -265,22 +265,22 @@ private:
265 }265 }
266 266 
267 // 从 cache 取出整行结果写入 rstdLocal[rowIndex]267 // 从 cache 取出整行结果写入 rstdLocal[rowIndex]
268- __local_mem__ float* dstPtr = (__local_mem__ float*)rstdLocal.GetPhyAddr();268+ __ubuf__ float* dstPtr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
269- __local_mem__ float* cachePtr = (__local_mem__ float*)cacheLocal.GetPhyAddr() +269+ __ubuf__ float* cachePtr = (__ubuf__ float*)cacheLocal.GetPhyAddr() +
270- tilingData_->resultCacheId * UB_BLOCK_SIZE_FP32;270+ tilingData_->resultCacheId * UB_BLOCK_SIZE_FP32;
271 __VEC_SCOPE__271 __VEC_SCOPE__
272 {272 {
273 RegTensor<float> a;273 RegTensor<float> a;
274 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();274 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
275- DataCopy<float, LoadDist::DIST_NORM>(a, cachePtr);275+ LoadAlign<float, LoadDist::DIST_NORM>(a, cachePtr);
276- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstPtr + rowIndex, a, pregOne);276+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstPtr + rowIndex, a, pregOne);
277 }277 }
278 }278 }
279 279 
280 __aicore__ inline void CastAndSquare(LocalTensor<T_X>& xLocal, LocalTensor<float>& xFp32Tmp, uint32_t count)280 __aicore__ inline void CastAndSquare(LocalTensor<T_X>& xLocal, LocalTensor<float>& xFp32Tmp, uint32_t count)
281 {281 {
282- __local_mem__ T_X* xAddr = (__local_mem__ T_X*)xLocal.GetPhyAddr();282+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
283- __local_mem__ float* dstAddr = (__local_mem__ float*)xFp32Tmp.GetPhyAddr();283+ __ubuf__ float* dstAddr = (__ubuf__ float*)xFp32Tmp.GetPhyAddr();
284 uint16_t loops = static_cast<uint16_t>(ops::CeilDiv(count, VL_FP32));284 uint16_t loops = static_cast<uint16_t>(ops::CeilDiv(count, VL_FP32));
285 285 
286 __VEC_SCOPE__286 __VEC_SCOPE__
@@ -293,7 +293,7 @@ private:
293 pregMask = UpdateMask<float>(sreg);293 pregMask = UpdateMask<float>(sreg);
294 LoadTensorForDtypeT(xAddr, xReg, pregMask, i * VL_FP32);294 LoadTensorForDtypeT(xAddr, xReg, pregMask, i * VL_FP32);
295 Mul(xReg, xReg, xReg, pregMask);295 Mul(xReg, xReg, xReg, pregMask);
296- DataCopy<float, StoreDist::DIST_NORM_B32>(dstAddr + i * VL_FP32, xReg, pregMask);296+ StoreAlign<float, StoreDist::DIST_NORM_B32>(dstAddr + i * VL_FP32, xReg, pregMask);
297 }297 }
298 }298 }
299 }299 }
@@ -301,9 +301,9 @@ private:
301 __aicore__ inline void FoldBlockVF(LocalTensor<T_X>& xLocal, LocalTensor<T_X>& xFoldLocal,301 __aicore__ inline void FoldBlockVF(LocalTensor<T_X>& xLocal, LocalTensor<T_X>& xFoldLocal,
302 LocalTensor<float>& xFp32Tmp, uint32_t tailCount, uint32_t count)302 LocalTensor<float>& xFp32Tmp, uint32_t tailCount, uint32_t count)
303 {303 {
304- __local_mem__ T_X* xInUb = (__local_mem__ T_X*)xLocal.GetPhyAddr();304+ __ubuf__ T_X* xInUb = (__ubuf__ T_X*)xLocal.GetPhyAddr();
305- __local_mem__ T_X* xFoldInUb = (__local_mem__ T_X*)xFoldLocal.GetPhyAddr();305+ __ubuf__ T_X* xFoldInUb = (__ubuf__ T_X*)xFoldLocal.GetPhyAddr();
306- __local_mem__ float* dstBuf = (__local_mem__ float*)xFp32Tmp.GetPhyAddr();306+ __ubuf__ float* dstBuf = (__ubuf__ float*)xFp32Tmp.GetPhyAddr();
307 307 
308 uint16_t loops = static_cast<uint16_t>(ops::CeilDiv(count, VL_FP32));308 uint16_t loops = static_cast<uint16_t>(ops::CeilDiv(count, VL_FP32));
309 uint16_t tailLoops = static_cast<uint16_t>(ops::CeilDiv(tailCount, VL_FP32));309 uint16_t tailLoops = static_cast<uint16_t>(ops::CeilDiv(tailCount, VL_FP32));
@@ -324,14 +324,14 @@ private:
324 Mul(xFoldReg, xFoldReg, xFoldReg, pregLoop);324 Mul(xFoldReg, xFoldReg, xFoldReg, pregLoop);
325 Add(sumReg, xReg, xFoldReg, pregLoop);325 Add(sumReg, xReg, xFoldReg, pregLoop);
326 Select(sumReg, sumReg, xReg, pregLoop); // 超出尾块范围用 xReg^2326 Select(sumReg, sumReg, xReg, pregLoop); // 超出尾块范围用 xReg^2
327- DataCopy<float, StoreDist::DIST_NORM_B32>(dstBuf + offset, sumReg, pregFull);327+ StoreAlign<float, StoreDist::DIST_NORM_B32>(dstBuf + offset, sumReg, pregFull);
328 }328 }
329 // 无尾块的部分：只有 x^2329 // 无尾块的部分：只有 x^2
330 for (uint16_t i = tailLoops; i < loops; ++i) {330 for (uint16_t i = tailLoops; i < loops; ++i) {
331 uint32_t offset = i * VL_FP32;331 uint32_t offset = i * VL_FP32;
332 LoadTensorForDtypeT(xInUb, xReg, pregFull, offset);332 LoadTensorForDtypeT(xInUb, xReg, pregFull, offset);
333 Mul(xReg, xReg, xReg, pregFull);333 Mul(xReg, xReg, xReg, pregFull);
334- DataCopy<float, StoreDist::DIST_NORM_B32>(dstBuf + offset, xReg, pregFull);334+ StoreAlign<float, StoreDist::DIST_NORM_B32>(dstBuf + offset, xReg, pregFull);
335 }335 }
336 }336 }
337 }337 }
@@ -344,21 +344,21 @@ private:
344 {344 {
345 uint16_t innerLoopTimes = cacheId;345 uint16_t innerLoopTimes = cacheId;
346 uint32_t innerLoopStride = stride;346 uint32_t innerLoopStride = stride;
347- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();347+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
348- __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride;348+ __ubuf__ float* cache = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheId * stride;
349- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();349+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
350 350 
351 __VEC_SCOPE__351 __VEC_SCOPE__
352 {352 {
353 RegTensor<float> aReg, bReg;353 RegTensor<float> aReg, bReg;
354 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();354 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
355 355 
356- DataCopy(aReg, (__local_mem__ float*)src);356+ LoadAlign(aReg, (__ubuf__ float*)src);
357 for (uint16_t j = 0; j < innerLoopTimes; ++j) {357 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
358- DataCopy(bReg, dst + j * innerLoopStride);358+ LoadAlign(bReg, dst + j * innerLoopStride);
359 Add(aReg, aReg, bReg, pregOne);359 Add(aReg, aReg, bReg, pregOne);
360 }360 }
361- DataCopy((__local_mem__ float*)cache, aReg, pregOne);361+ StoreAlign((__ubuf__ float*)cache, aReg, pregOne);
362 }362 }
363 }363 }
364 364 
@@ -415,14 +415,14 @@ private:
415 LocalTensor<T_GAMMA>& betaLocal, LocalTensor<float>& rstdLocal,415 LocalTensor<T_GAMMA>& betaLocal, LocalTensor<float>& rstdLocal,
416 LocalTensor<T_X>& yLocal, uint32_t rstdOffset, uint32_t curN)416 LocalTensor<T_X>& yLocal, uint32_t rstdOffset, uint32_t curN)
417 {417 {
418- __local_mem__ T_X* xAddr = (__local_mem__ T_X*)xLocal.GetPhyAddr();418+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
419- __local_mem__ T_GAMMA* gammaAddr = (__local_mem__ T_GAMMA*)gammaLocal.GetPhyAddr();419+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
420- __local_mem__ T_GAMMA* betaAddr;420+ __ubuf__ T_GAMMA* betaAddr;
421 if constexpr (hasInputBeta) {421 if constexpr (hasInputBeta) {
422- betaAddr = (__local_mem__ T_GAMMA*)betaLocal.GetPhyAddr();422+ betaAddr = (__ubuf__ T_GAMMA*)betaLocal.GetPhyAddr();
423 }423 }
424- __local_mem__ float* rstdAddr = (__local_mem__ float*)rstdLocal.GetPhyAddr();424+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
425- __local_mem__ T_X* yAddr = (__local_mem__ T_X*)yLocal.GetPhyAddr();425+ __ubuf__ T_X* yAddr = (__ubuf__ T_X*)yLocal.GetPhyAddr();
426 426 
427 uint16_t nloops = static_cast<uint16_t>(427 uint16_t nloops = static_cast<uint16_t>(
428 ops::CeilDiv(static_cast<uint64_t>(curN), static_cast<uint64_t>(VL_FP32)));428 ops::CeilDiv(static_cast<uint64_t>(curN), static_cast<uint64_t>(VL_FP32)));
@@ -438,7 +438,7 @@ private:
438 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();438 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
439 439 
440 uint32_t sreg = curN;440 uint32_t sreg = curN;
441- DataCopy<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdAddr + rstdOffset);441+ LoadAlign<float, LoadDist::DIST_BRC_B32>(RstdReg, rstdAddr + rstdOffset);
442 for (uint16_t j = 0; j < nloops; ++j) {442 for (uint16_t j = 0; j < nloops; ++j) {
443 pregMask = UpdateMask<float>(sreg);443 pregMask = UpdateMask<float>(sreg);
444 uint32_t off = j * VL_FP32;444 uint32_t off = j * VL_FP32;
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_dgamma_regbase_helper.h+66-67
@@ -44,75 +44,74 @@ __aicore__ inline uint32_t BLOCK_ALIGN(uint32_t x, uint32_t blockSize)
44}44}
45 45 
46template <typename DY_TYPE, typename X_TYPE, typename RSTD_TYPE, int TILING_KEY>46template <typename DY_TYPE, typename X_TYPE, typename RSTD_TYPE, int TILING_KEY>
47-__aicore__ inline void CalcMulRes(__local_mem__ DY_TYPE* dyAddr, __local_mem__ X_TYPE* xAddr,47+__aicore__ inline void CalcMulRes(__ubuf__ DY_TYPE* dyAddr, __ubuf__ X_TYPE* xAddr, __ubuf__ RSTD_TYPE* rstdAddr,
48- __local_mem__ RSTD_TYPE* rstdAddr, __local_mem__ float* dgammaOutAddr, MaskReg& preg,48+ __ubuf__ float* dgammaOutAddr, MaskReg& preg, uint32_t offset0, uint32_t k)
49- uint32_t offset0, uint32_t k)
50{49{
51 RegTensor<float> xFp32, rstdFp32, dyFp32, temp_res, mul_res;50 RegTensor<float> xFp32, rstdFp32, dyFp32, temp_res, mul_res;
52 51 
53 if constexpr (IsSameType<DY_TYPE, float>::value) {52 if constexpr (IsSameType<DY_TYPE, float>::value) {
54- DataCopy<DY_TYPE, LoadDist::DIST_NORM>(dyFp32, (__local_mem__ float*)(dyAddr + offset0));53+ LoadAlign<DY_TYPE, LoadDist::DIST_NORM>(dyFp32, (__ubuf__ float*)(dyAddr + offset0));
55 } else {54 } else {
56 RegTensor<DY_TYPE> dstRegB16;55 RegTensor<DY_TYPE> dstRegB16;
57- DataCopy<DY_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__local_mem__ DY_TYPE*)(dyAddr + offset0));56+ LoadAlign<DY_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__ubuf__ DY_TYPE*)(dyAddr + offset0));
58 Cast<float, DY_TYPE, castTraitB162B32>(dyFp32, dstRegB16, preg);57 Cast<float, DY_TYPE, castTraitB162B32>(dyFp32, dstRegB16, preg);
59 }58 }
60 59 
61 if constexpr (IsSameType<X_TYPE, float>::value) {60 if constexpr (IsSameType<X_TYPE, float>::value) {
62- DataCopy<X_TYPE, LoadDist::DIST_NORM>(xFp32, (__local_mem__ float*)(xAddr + offset0));61+ LoadAlign<X_TYPE, LoadDist::DIST_NORM>(xFp32, (__ubuf__ float*)(xAddr + offset0));
63 } else {62 } else {
64 RegTensor<X_TYPE> dstRegB16;63 RegTensor<X_TYPE> dstRegB16;
65- DataCopy<X_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__local_mem__ X_TYPE*)(xAddr + offset0));64+ LoadAlign<X_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__ubuf__ X_TYPE*)(xAddr + offset0));
66 Cast<float, X_TYPE, castTraitB162B32>(xFp32, dstRegB16, preg);65 Cast<float, X_TYPE, castTraitB162B32>(xFp32, dstRegB16, preg);
67 }66 }
68 67 
69- DataCopy<RSTD_TYPE, LoadDist::DIST_BRC_B32>(rstdFp32, ((__local_mem__ float*)rstdAddr + k));68+ LoadAlign<RSTD_TYPE, LoadDist::DIST_BRC_B32>(rstdFp32, ((__ubuf__ float*)rstdAddr + k));
70 69 
71 Mul(temp_res, xFp32, rstdFp32, preg);70 Mul(temp_res, xFp32, rstdFp32, preg);
72 Mul(mul_res, dyFp32, temp_res, preg);71 Mul(mul_res, dyFp32, temp_res, preg);
73 72 
74- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dgammaOutAddr + offset0), mul_res, preg);73+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dgammaOutAddr + offset0), mul_res, preg);
75}74}
76 75 
77-__aicore__ inline void reduceSumCompressedBy8(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset,76+__aicore__ inline void reduceSumCompressedBy8(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset,
78 uint32_t ub_offset)77 uint32_t ub_offset)
79{78{
80 RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,79 RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,
81 temp_reg3_1, temp_reg4_0, temp_reg4_1, temp_reg5_0, temp_reg5_1, temp_reg6_0, temp_reg6_1, temp_reg7_0,80 temp_reg3_1, temp_reg4_0, temp_reg4_1, temp_reg5_0, temp_reg5_1, temp_reg6_0, temp_reg6_1, temp_reg7_0,
82 temp_reg7_1;81 temp_reg7_1;
83- __local_mem__ float* currentAddr = dyAddr + REDUCEBY8ELENUM * ub_offset;82+ __ubuf__ float* currentAddr = dyAddr + REDUCEBY8ELENUM * ub_offset;
84 83 
85 //84 //
86- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(currentAddr));85+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(currentAddr));
87- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(currentAddr + offset));86+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(currentAddr + offset));
88 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);87 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
89 88 
90- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_0, (__local_mem__ float*)(currentAddr + 2 * offset));89+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_0, (__ubuf__ float*)(currentAddr + 2 * offset));
91- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_1, (__local_mem__ float*)(currentAddr + 3 * offset));90+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_1, (__ubuf__ float*)(currentAddr + 3 * offset));
92 AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);91 AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);
93 92 
94- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_0, (__local_mem__ float*)(currentAddr + 4 * offset));93+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_0, (__ubuf__ float*)(currentAddr + 4 * offset));
95- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_1, (__local_mem__ float*)(currentAddr + 5 * offset));94+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_1, (__ubuf__ float*)(currentAddr + 5 * offset));
96 AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);95 AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);
97 96 
98- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_0, (__local_mem__ float*)(currentAddr + 6 * offset));97+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_0, (__ubuf__ float*)(currentAddr + 6 * offset));
99- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_1, (__local_mem__ float*)(currentAddr + 7 * offset));98+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_1, (__ubuf__ float*)(currentAddr + 7 * offset));
100 AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);99 AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);
101 100 
102- DataCopy<float, LoadDist::DIST_NORM>(temp_reg4_0, (__local_mem__ float*)(currentAddr + 8 * offset));101+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg4_0, (__ubuf__ float*)(currentAddr + 8 * offset));
103- DataCopy<float, LoadDist::DIST_NORM>(temp_reg4_1, (__local_mem__ float*)(currentAddr + 9 * offset));102+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg4_1, (__ubuf__ float*)(currentAddr + 9 * offset));
104 AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg4_1, preg);103 AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg4_1, preg);
105 104 
106- DataCopy<float, LoadDist::DIST_NORM>(temp_reg5_0, (__local_mem__ float*)(currentAddr + 10 * offset));105+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg5_0, (__ubuf__ float*)(currentAddr + 10 * offset));
107- DataCopy<float, LoadDist::DIST_NORM>(temp_reg5_1, (__local_mem__ float*)(currentAddr + 11 * offset));106+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg5_1, (__ubuf__ float*)(currentAddr + 11 * offset));
108 AscendC::MicroAPI::Add(temp_reg5_0, temp_reg5_0, temp_reg5_1, preg);107 AscendC::MicroAPI::Add(temp_reg5_0, temp_reg5_0, temp_reg5_1, preg);
109 108 
110- DataCopy<float, LoadDist::DIST_NORM>(temp_reg6_0, (__local_mem__ float*)(currentAddr + 12 * offset));109+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg6_0, (__ubuf__ float*)(currentAddr + 12 * offset));
111- DataCopy<float, LoadDist::DIST_NORM>(temp_reg6_1, (__local_mem__ float*)(currentAddr + 13 * offset));110+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg6_1, (__ubuf__ float*)(currentAddr + 13 * offset));
112 AscendC::MicroAPI::Add(temp_reg6_0, temp_reg6_0, temp_reg6_1, preg);111 AscendC::MicroAPI::Add(temp_reg6_0, temp_reg6_0, temp_reg6_1, preg);
113 112 
114- DataCopy<float, LoadDist::DIST_NORM>(temp_reg7_0, (__local_mem__ float*)(currentAddr + 14 * offset));113+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg7_0, (__ubuf__ float*)(currentAddr + 14 * offset));
115- DataCopy<float, LoadDist::DIST_NORM>(temp_reg7_1, (__local_mem__ float*)(currentAddr + 15 * offset));114+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg7_1, (__ubuf__ float*)(currentAddr + 15 * offset));
116 AscendC::MicroAPI::Add(temp_reg7_0, temp_reg7_0, temp_reg7_1, preg);115 AscendC::MicroAPI::Add(temp_reg7_0, temp_reg7_0, temp_reg7_1, preg);
117 116 
118 //117 //
@@ -126,29 +125,29 @@ __aicore__ inline void reduceSumCompressedBy8(__local_mem__ float* dyAddr, MaskR
126 125 
127 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg4_0, preg);126 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg4_0, preg);
128 127 
129- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);128+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);
130}129}
131 130 
132-__aicore__ inline void reduceSumCompressedBy4(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset,131+__aicore__ inline void reduceSumCompressedBy4(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset,
133 uint32_t ub_offset)132 uint32_t ub_offset)
134{133{
135 RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,134 RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,
136 temp_reg3_1;135 temp_reg3_1;
137- __local_mem__ float* currentAddr = dyAddr + COMPRESSBY8ELENUM * ub_offset;136+ __ubuf__ float* currentAddr = dyAddr + COMPRESSBY8ELENUM * ub_offset;
138- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(currentAddr));137+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(currentAddr));
139- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(currentAddr + offset));138+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(currentAddr + offset));
140 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);139 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
141 140 
142- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_0, (__local_mem__ float*)(currentAddr + 2 * offset));141+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_0, (__ubuf__ float*)(currentAddr + 2 * offset));
143- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_1, (__local_mem__ float*)(currentAddr + 3 * offset));142+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_1, (__ubuf__ float*)(currentAddr + 3 * offset));
144 AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);143 AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);
145 144 
146- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_0, (__local_mem__ float*)(currentAddr + 4 * offset));145+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_0, (__ubuf__ float*)(currentAddr + 4 * offset));
147- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_1, (__local_mem__ float*)(currentAddr + 5 * offset));146+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_1, (__ubuf__ float*)(currentAddr + 5 * offset));
148 AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);147 AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);
149 148 
150- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_0, (__local_mem__ float*)(currentAddr + 6 * offset));149+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_0, (__ubuf__ float*)(currentAddr + 6 * offset));
151- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_1, (__local_mem__ float*)(currentAddr + 7 * offset));150+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_1, (__ubuf__ float*)(currentAddr + 7 * offset));
152 AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);151 AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);
153 152 
154 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);153 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);
@@ -156,69 +155,69 @@ __aicore__ inline void reduceSumCompressedBy4(__local_mem__ float* dyAddr, MaskR
156 155 
157 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg2_0, preg);156 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg2_0, preg);
158 157 
159- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);158+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);
160}159}
161 160 
162-__aicore__ inline void reduceSumCompressedBy2(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset,161+__aicore__ inline void reduceSumCompressedBy2(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset,
163 uint32_t ub_offset)162 uint32_t ub_offset)
164{163{
165 RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1;164 RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1;
166 165 
167- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(dyAddr));166+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(dyAddr));
168- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(dyAddr + offset));167+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(dyAddr + offset));
169 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);168 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
170 169 
171- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_0, (__local_mem__ float*)(dyAddr + 2 * offset));170+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_0, (__ubuf__ float*)(dyAddr + 2 * offset));
172- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_1, (__local_mem__ float*)(dyAddr + 3 * offset));171+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_1, (__ubuf__ float*)(dyAddr + 3 * offset));
173 AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);172 AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);
174 173 
175 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);174 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);
176- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);175+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);
177}176}
178 177 
179-__aicore__ inline void reduceSumCompressedBy1(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset)178+__aicore__ inline void reduceSumCompressedBy1(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset)
180{179{
181 RegTensor<float> temp_reg0_0, temp_reg0_1;180 RegTensor<float> temp_reg0_0, temp_reg0_1;
182 181 
183- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(dyAddr));182+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(dyAddr));
184- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(dyAddr + offset));183+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(dyAddr + offset));
185 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);184 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
186 185 
187- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr), temp_reg0_0, preg);186+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr), temp_reg0_0, preg);
188}187}
189 188 
190-__aicore__ inline void reduceSumCompressedBy8WithOutPad(__local_mem__ float* src1Addr, __local_mem__ float* src2Addr,189+__aicore__ inline void reduceSumCompressedBy8WithOutPad(__ubuf__ float* src1Addr, __ubuf__ float* src2Addr,
191 MaskReg& preg, uint32_t ub_offset, uint32_t vlFp32)190 MaskReg& preg, uint32_t ub_offset, uint32_t vlFp32)
192{191{
193 for (uint16_t i = 0; i < 8; i++) {192 for (uint16_t i = 0; i < 8; i++) {
194 RegTensor<float> temp_reg0_0, temp_reg0_1;193 RegTensor<float> temp_reg0_0, temp_reg0_1;
195 uint32_t tempOffset = i * vlFp32;194 uint32_t tempOffset = i * vlFp32;
196- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(src1Addr + ub_offset + tempOffset));195+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(src1Addr + ub_offset + tempOffset));
197- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(src2Addr + ub_offset + tempOffset));196+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(src2Addr + ub_offset + tempOffset));
198 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);197 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
199- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(src1Addr + ub_offset + tempOffset),198+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(src1Addr + ub_offset + tempOffset), temp_reg0_0,
200- temp_reg0_0, preg);199+ preg);
201 }200 }
202}201}
203 202 
204-__aicore__ inline void reduceSumCompressedBy8WithPad(__local_mem__ float* src1Addr, __local_mem__ float* src2Addr,203+__aicore__ inline void reduceSumCompressedBy8WithPad(__ubuf__ float* src1Addr, __ubuf__ float* src2Addr, MaskReg& preg,
205- MaskReg& preg, uint32_t ub_offset, uint32_t rowsBoundLine,204+ uint32_t ub_offset, uint32_t rowsBoundLine, uint32_t vlFp32,
206- uint32_t vlFp32, uint32_t tailDataOffset)205+ uint32_t tailDataOffset)
207{206{
208 for (uint16_t i = 0; i < 8; i++) {207 for (uint16_t i = 0; i < 8; i++) {
209 RegTensor<float> temp_reg0_0, temp_reg0_1;208 RegTensor<float> temp_reg0_0, temp_reg0_1;
210 uint32_t temp_off_set_0 = ub_offset + i * vlFp32;209 uint32_t temp_off_set_0 = ub_offset + i * vlFp32;
211 uint32_t temp_off_set_1 = tailDataOffset + temp_off_set_0 < rowsBoundLine ? tailDataOffset + temp_off_set_0 :210 uint32_t temp_off_set_1 = tailDataOffset + temp_off_set_0 < rowsBoundLine ? tailDataOffset + temp_off_set_0 :
212 rowsBoundLine;211 rowsBoundLine;
213- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(src1Addr + temp_off_set_0));212+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(src1Addr + temp_off_set_0));
214- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(src2Addr + temp_off_set_1));213+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(src2Addr + temp_off_set_1));
215 214 
216 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);215 AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
217- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(src1Addr + temp_off_set_0), temp_reg0_0, preg);216+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(src1Addr + temp_off_set_0), temp_reg0_0, preg);
218 }217 }
219}218}
220 219 
221-__aicore__ inline void UpdateCache(const AscendC::LocalTensor<float>& dstTensor, __local_mem__ float* srcAddr,220+__aicore__ inline void UpdateCache(const AscendC::LocalTensor<float>& dstTensor, __ubuf__ float* srcAddr,
222 const int64_t cacheID, const int64_t count)221 const int64_t cacheID, const int64_t count)
223{222{
224 // UpdateCache223 // UpdateCache
@@ -226,21 +225,21 @@ __aicore__ inline void UpdateCache(const AscendC::LocalTensor<float>& dstTensor,
226 uint32_t innerLoopStride = count;225 uint32_t innerLoopStride = count;
227 __VEC_SCOPE__226 __VEC_SCOPE__
228 {227 {
229- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();228+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
230- __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * count;229+ __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * count;
231 uint32_t sreg = static_cast<uint32_t>(count);230 uint32_t sreg = static_cast<uint32_t>(count);
232 AscendC::MicroAPI::RegTensor<float> aReg, bReg;231 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
233 AscendC::MicroAPI::MaskReg pMask;232 AscendC::MicroAPI::MaskReg pMask;
234 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);233 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
235- DataCopy(aReg, (__local_mem__ float*)srcAddr);234+ LoadAlign(aReg, (__ubuf__ float*)srcAddr);
236 for (uint16_t j = 0; j < innerLoopTimes; ++j) {235 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
237- DataCopy(bReg, (__local_mem__ float*)dst + static_cast<uint32_t>(j * innerLoopStride));236+ LoadAlign(bReg, (__ubuf__ float*)dst + static_cast<uint32_t>(j * innerLoopStride));
238 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);237 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
239 }238 }
240- DataCopy((__local_mem__ float*)cah, aReg, pMask);239+ StoreAlign((__ubuf__ float*)cah, aReg, pMask);
241 }240 }
242}241}
243 242 
244__aicore__ inline int64_t GetCacheID(const int64_t idx) { return ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1; }243__aicore__ inline int64_t GetCacheID(const int64_t idx) { return ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1; }
245} // namespace RmsNormGrad244} // namespace RmsNormGrad
246-#endif // RMS_NORM_GRAD_REGBASE_DGAMMA_H245+#endif // RMS_NORM_GRAD_REGBASE_DGAMMA_H
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_regbase_common.h+69-71
@@ -102,15 +102,14 @@ __aicore__ inline int32_t findPowerTwo(int32_t n)
102}102}
103 103 
104template <typename T>104template <typename T>
105-__aicore__ inline void LoadAndCast(RegTensor<float>& dstReg, __local_mem__ T* srcAddr, MaskReg& maskReg,105+__aicore__ inline void LoadAndCast(RegTensor<float>& dstReg, __ubuf__ T* srcAddr, MaskReg& maskReg, uint32_t srcOffset)
106- uint32_t srcOffset)
107{106{
108 if constexpr (IsSameType<T, float>::value) {107 if constexpr (IsSameType<T, float>::value) {
109- DataCopy(dstReg, srcAddr + srcOffset);108+ LoadAlign(dstReg, srcAddr + srcOffset);
110 } else {109 } else {
111 RegTensor<T> dstRegB16;110 RegTensor<T> dstRegB16;
112- // DataCopy<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + ubFactorD_ + srcOffset); // 后续不做偏移111+ // LoadAlign<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + ubFactorD_ + srcOffset); // 后续不做偏移
113- DataCopy<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + srcOffset); // 后续不做偏移112+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + srcOffset); // 后续不做偏移
114 Cast<float, T, castTraitB162B32>(dstReg, dstRegB16, maskReg);113 Cast<float, T, castTraitB162B32>(dstReg, dstRegB16, maskReg);
115 }114 }
116}115}
@@ -132,11 +131,11 @@ __aicore__ inline void LevelMerge(LocalTensor<float>& dstLocal, LocalTensor<floa
132 uint32_t sreg = (uint32_t)(calCount);131 uint32_t sreg = (uint32_t)(calCount);
133 uint32_t meanTile = repeatTimes;132 uint32_t meanTile = repeatTimes;
134 133 
135- __local_mem__ float* src1Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 0 * calCount;134+ __ubuf__ float* src1Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 0 * calCount;
136- __local_mem__ float* src2Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 1 * calCount;135+ __ubuf__ float* src2Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 1 * calCount;
137- __local_mem__ float* src3Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 2 * calCount;136+ __ubuf__ float* src3Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 2 * calCount;
138- __local_mem__ float* src4Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 3 * calCount;137+ __ubuf__ float* src4Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 3 * calCount;
139- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();138+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
140 139 
141 __VEC_SCOPE__140 __VEC_SCOPE__
142 {141 {
@@ -145,16 +144,16 @@ __aicore__ inline void LevelMerge(LocalTensor<float>& dstLocal, LocalTensor<floa
145 MaskReg pregLoop;144 MaskReg pregLoop;
146 for (uint16_t i = 0; i < repeatTimes; ++i) {145 for (uint16_t i = 0; i < repeatTimes; ++i) {
147 pregLoop = UpdateMask<float>(sreg);146 pregLoop = UpdateMask<float>(sreg);
148- DataCopy(vRegA, src1Addr + static_cast<uint32_t>(i * V_LENGTH));147+ LoadAlign(vRegA, src1Addr + static_cast<uint32_t>(i * V_LENGTH));
149- DataCopy(vRegB, src2Addr + static_cast<uint32_t>(i * V_LENGTH));148+ LoadAlign(vRegB, src2Addr + static_cast<uint32_t>(i * V_LENGTH));
150- DataCopy(vRegC, src3Addr + static_cast<uint32_t>(i * V_LENGTH));149+ LoadAlign(vRegC, src3Addr + static_cast<uint32_t>(i * V_LENGTH));
151- DataCopy(vRegD, src4Addr + static_cast<uint32_t>(i * V_LENGTH));150+ LoadAlign(vRegD, src4Addr + static_cast<uint32_t>(i * V_LENGTH));
152 Add(vRegA, vRegA, vRegB, pregLoop);151 Add(vRegA, vRegA, vRegB, pregLoop);
153 Add(vRegC, vRegC, vRegD, pregLoop);152 Add(vRegC, vRegC, vRegD, pregLoop);
154 Add(dstReg, vRegA, vRegC, pregLoop);153 Add(dstReg, vRegA, vRegC, pregLoop);
155- ReduceSum(vMean, dstReg, pregLoop);154+ Reduce<ReduceType::SUM>(vMean, dstReg, pregLoop);
156- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(offset), vMean,155+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(offset), vMean,
157- pregMerge);156+ pregMerge);
158 }157 }
159 }158 }
160}159}
@@ -239,11 +238,11 @@ __aicore__ inline void ReduceSumImpl(LocalTensor<float>& dstLocal, LocalTensor<f
239 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;238 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;
240 uint32_t meanSreg = meanTile;239 uint32_t meanSreg = meanTile;
241 240 
242- __local_mem__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();241+ __ubuf__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();
243- __local_mem__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(powerSplit);242+ __ubuf__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(powerSplit);
244- __local_mem__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(remainTile);243+ __ubuf__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(remainTile);
245- __local_mem__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();244+ __ubuf__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();
246- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();245+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
247 246 
248 __VEC_SCOPE__247 __VEC_SCOPE__
249 {248 {
@@ -253,43 +252,43 @@ __aicore__ inline void ReduceSumImpl(LocalTensor<float>& dstLocal, LocalTensor<f
253 252 
254 for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {253 for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {
255 pregLoop = UpdateMask<float>(remainSreg);254 pregLoop = UpdateMask<float>(remainSreg);
256- DataCopy(mainAReg, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));255+ LoadAlign(mainAReg, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
257- DataCopy(mainBReg, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));256+ LoadAlign(mainBReg, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
258- DataCopy(tailAReg, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));257+ LoadAlign(tailAReg, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
259- DataCopy(tailBReg, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));258+ LoadAlign(tailBReg, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
260 259 
261 Add(mainAReg, mainAReg, tailAReg, pregLoop);260 Add(mainAReg, mainAReg, tailAReg, pregLoop);
262 Add(mainBReg, mainBReg, tailBReg, pregLoop);261 Add(mainBReg, mainBReg, tailBReg, pregLoop);
263 Add(mainAReg, mainAReg, mainBReg, pregLoop);262 Add(mainAReg, mainAReg, mainBReg, pregLoop);
264- ReduceSum(vMeanReg, mainAReg, pregLoop);263+ Reduce<ReduceType::SUM>(vMeanReg, mainAReg, pregLoop);
265- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMeanReg,264+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMeanReg,
266- pregMerge);265+ pregMerge);
267 }266 }
268 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {267 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
269 pregLoop = UpdateMask<float>(masterSreg);268 pregLoop = UpdateMask<float>(masterSreg);
270- DataCopy(mainAReg, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));269+ LoadAlign(mainAReg, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
271- DataCopy(mainBReg, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));270+ LoadAlign(mainBReg, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
272 Add(mainAReg, mainAReg, mainBReg, pregLoop);271 Add(mainAReg, mainAReg, mainBReg, pregLoop);
273- ReduceSum(vMeanReg, mainAReg, pregLoop);272+ Reduce<ReduceType::SUM>(vMeanReg, mainAReg, pregLoop);
274- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(remainRepeats + i),273+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(remainRepeats + i),
275- vMeanReg, pregMerge);274+ vMeanReg, pregMerge);
276 }275 }
277 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();276 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
278 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {277 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
279 pregLoop = UpdateMask<float>(mergeSreg);278 pregLoop = UpdateMask<float>(mergeSreg);
280- DataCopy(mainAReg, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));279+ LoadAlign(mainAReg, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
281- DataCopy(mainBReg, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));280+ LoadAlign(mainBReg, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
282 Add(mainAReg, mainAReg, mainBReg, pregLoop);281 Add(mainAReg, mainAReg, mainBReg, pregLoop);
283- ReduceSum(vMeanReg, mainAReg, pregLoop);282+ Reduce<ReduceType::SUM>(vMeanReg, mainAReg, pregLoop);
284- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMeanReg,283+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMeanReg,
285- pregMerge);284+ pregMerge);
286 }285 }
287 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();286 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
288 {287 {
289 pregLoop = UpdateMask<float>(meanSreg);288 pregLoop = UpdateMask<float>(meanSreg);
290- DataCopy(mainAReg, workAddr + 0);289+ LoadAlign(mainAReg, workAddr + 0);
291- ReduceSum(vMeanReg, mainAReg, pregLoop);290+ Reduce<ReduceType::SUM>(vMeanReg, mainAReg, pregLoop);
292- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMeanReg, pregMerge);291+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMeanReg, pregMerge);
293 }292 }
294 }293 }
295}294}
@@ -319,8 +318,8 @@ __aicore__ inline void MultiReduceSumImpl(LocalTensor<float>& dstLocal, LocalTen
319 318 
320 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;319 uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;
321 320 
322- __local_mem__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();321+ __ubuf__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();
323- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();322+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
324 323 
325 __VEC_SCOPE__324 __VEC_SCOPE__
326 {325 {
@@ -331,59 +330,58 @@ __aicore__ inline void MultiReduceSumImpl(LocalTensor<float>& dstLocal, LocalTen
331 uint32_t mergeSreg = mergeTile;330 uint32_t mergeSreg = mergeTile;
332 uint32_t meanSreg = meanTile;331 uint32_t meanSreg = meanTile;
333 332 
334- __local_mem__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL;333+ __ubuf__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL;
335- __local_mem__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL +334+ __ubuf__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL + int64_t(powerSplit);
336- int64_t(powerSplit);335+ __ubuf__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL +
337- __local_mem__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL +336+ int64_t(remainTile);
338- int64_t(remainTile);
339 337 
340 MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();338 MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();
341 MaskReg pregMask;339 MaskReg pregMask;
342 340 
343 for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {341 for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {
344 pregMask = UpdateMask<float>(remainSreg);342 pregMask = UpdateMask<float>(remainSreg);
345- DataCopy(mainA, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));343+ LoadAlign(mainA, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
346- DataCopy(mainB, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));344+ LoadAlign(mainB, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
347- DataCopy(tailA, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));345+ LoadAlign(tailA, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
348- DataCopy(tailB, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));346+ LoadAlign(tailB, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
349 347 
350 Add(mainA, mainA, tailA, pregMask);348 Add(mainA, mainA, tailA, pregMask);
351 Add(mainB, mainB, tailB, pregMask);349 Add(mainB, mainB, tailB, pregMask);
352 Add(mainA, mainA, mainB, pregMask);350 Add(mainA, mainA, mainB, pregMask);
353- ReduceSum(vMean, mainA, pregMask);351+ Reduce<ReduceType::SUM>(vMean, mainA, pregMask);
354- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,352+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,
355- pregMerge);353+ pregMerge);
356 }354 }
357 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {355 for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
358 pregMask = UpdateMask<float>(masterSreg);356 pregMask = UpdateMask<float>(masterSreg);
359- DataCopy(mainA, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));357+ LoadAlign(mainA, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
360- DataCopy(mainB, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));358+ LoadAlign(mainB, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
361 Add(mainA, mainA, mainB, pregMask);359 Add(mainA, mainA, mainB, pregMask);
362- ReduceSum(vMean, mainA, pregMask);360+ Reduce<ReduceType::SUM>(vMean, mainA, pregMask);
363- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(remainRepeats + i),361+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
364- vMean, pregMerge);362+ workAddr + static_cast<uint32_t>(remainRepeats + i), vMean, pregMerge);
365 }363 }
366 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();364 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
367 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {365 for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
368 pregMask = UpdateMask<float>(mergeSreg);366 pregMask = UpdateMask<float>(mergeSreg);
369- DataCopy(mainA, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));367+ LoadAlign(mainA, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
370- DataCopy(mainB, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));368+ LoadAlign(mainB, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
371 Add(mainA, mainA, mainB, pregMask);369 Add(mainA, mainA, mainB, pregMask);
372- ReduceSum(vMean, mainA, pregMask);370+ Reduce<ReduceType::SUM>(vMean, mainA, pregMask);
373- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,371+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,
374- pregMerge);372+ pregMerge);
375 }373 }
376 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();374 LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
377 {375 {
378 pregMask = UpdateMask<float>(meanSreg);376 pregMask = UpdateMask<float>(meanSreg);
379- DataCopy(mainA, workAddr + 0);377+ LoadAlign(mainA, workAddr + 0);
380- ReduceSum(vMean, mainA, pregMask);378+ Reduce<ReduceType::SUM>(vMean, mainA, pregMask);
381- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(r), vMean,379+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(r), vMean,
382- pregMerge);380+ pregMerge);
383 }381 }
384 }382 }
385 }383 }
386}384}
387 385 
388} // namespace RmsNormGrad386} // namespace RmsNormGrad
389-#endif // OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_REGBASE_COMMON_H387+#endif // OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_REGBASE_COMMON_H
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_regbase_dgamma.h+24-25
@@ -103,8 +103,8 @@ public:
103 DataCopyPad(dgammaGm_[dgammaGmOffset], outLocal[dgammaUBOffset], dataCopyParams);103 DataCopyPad(dgammaGm_[dgammaGmOffset], outLocal[dgammaUBOffset], dataCopyParams);
104 }104 }
105 105 
106- __aicore__ inline void VFCalcPreDgamma(__local_mem__ DY_TYPE* dyAddr, __local_mem__ X_TYPE* xAddr,106+ __aicore__ inline void VFCalcPreDgamma(__ubuf__ DY_TYPE* dyAddr, __ubuf__ X_TYPE* xAddr,
107- __local_mem__ RSTD_TYPE* rstdAddr, __local_mem__ float* dgammaOutAddr,107+ __ubuf__ RSTD_TYPE* rstdAddr, __ubuf__ float* dgammaOutAddr,
108 uint16_t curUBLoopColsCount, int32_t curRowsNum)108 uint16_t curUBLoopColsCount, int32_t curRowsNum)
109 {109 {
110 uint16_t colsRegLoopCount = CEIL_DIV(curUBLoopColsCount, vlFp32_);110 uint16_t colsRegLoopCount = CEIL_DIV(curUBLoopColsCount, vlFp32_);
@@ -122,7 +122,7 @@ public:
122 }122 }
123 }123 }
124 124 
125- __aicore__ inline void VFDuplicateRows(__local_mem__ float* srcAddr, uint32_t padRowsLen, uint64_t rowsBoundLine)125+ __aicore__ inline void VFDuplicateRows(__ubuf__ float* srcAddr, uint32_t padRowsLen, uint64_t rowsBoundLine)
126 {126 {
127 __VEC_SCOPE__127 __VEC_SCOPE__
128 {128 {
@@ -131,11 +131,11 @@ public:
131 MaskReg pregLoop = UpdateMask<float>(sreg0);131 MaskReg pregLoop = UpdateMask<float>(sreg0);
132 // 填充数据132 // 填充数据
133 AscendC::MicroAPI::Duplicate(tempReg, 0);133 AscendC::MicroAPI::Duplicate(tempReg, 0);
134- AscendC::MicroAPI::DataCopy(srcAddr + static_cast<uint32_t>(rowsBoundLine), tempReg, pregLoop);134+ AscendC::MicroAPI::StoreAlign(srcAddr + static_cast<uint32_t>(rowsBoundLine), tempReg, pregLoop);
135 }135 }
136 }136 }
137 137 
138- __aicore__ inline void VFBinaryReduceSumWithoutTail(__local_mem__ float* dgammaOutAddr, uint16_t curUbLoopColsCount,138+ __aicore__ inline void VFBinaryReduceSumWithoutTail(__ubuf__ float* dgammaOutAddr, uint16_t curUbLoopColsCount,
139 int64_t rows)139 int64_t rows)
140 {140 {
141 uint32_t BinaryAddNumLevel2 = rows / REDUCEBY8ELENUM;141 uint32_t BinaryAddNumLevel2 = rows / REDUCEBY8ELENUM;
@@ -187,8 +187,7 @@ public:
187 }187 }
188 }188 }
189 189 
190- __aicore__ inline void VFHandleTailRows(__local_mem__ float* dgammaOutAddr, uint16_t rowsTail,190+ __aicore__ inline void VFHandleTailRows(__ubuf__ float* dgammaOutAddr, uint16_t rowsTail, uint64_t tailDataOffset)
191- uint64_t tailDataOffset)
192 {191 {
193 uint32_t BinaryAddTailNum = (rowsTail + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;192 uint32_t BinaryAddTailNum = (rowsTail + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;
194 193 
@@ -206,8 +205,8 @@ public:
206 }205 }
207 }206 }
208 207 
209- __aicore__ inline void VFHandleTailRowsWithTwoBuffer(__local_mem__ float* dgammaOutAddr,208+ __aicore__ inline void VFHandleTailRowsWithTwoBuffer(__ubuf__ float* dgammaOutAddr, __ubuf__ float* dgammaOutAddr1,
210- __local_mem__ float* dgammaOutAddr1, uint64_t tailRowsNum)209+ uint64_t tailRowsNum)
211 {210 {
212 uint32_t BinaryAddTailNum = (tailRowsNum + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;211 uint32_t BinaryAddTailNum = (tailRowsNum + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;
213 uint32_t tailDataOffset = 0;212 uint32_t tailDataOffset = 0;
@@ -242,10 +241,10 @@ public:
242 xLocal = xQueue_.template DeQue<X_TYPE>();241 xLocal = xQueue_.template DeQue<X_TYPE>();
243 rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();242 rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
244 243 
245- __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();244+ __ubuf__ DY_TYPE* dyAddr = (__ubuf__ DY_TYPE*)dyLocal[0].GetPhyAddr();
246- __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();245+ __ubuf__ X_TYPE* xAddr = (__ubuf__ X_TYPE*)xLocal[0].GetPhyAddr();
247- __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();246+ __ubuf__ RSTD_TYPE* rstdAddr = (__ubuf__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
248- __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();247+ __ubuf__ float* dgammaOutAddr = (__ubuf__ float*)dgammaOutLocal[0].GetPhyAddr();
249 248 
250 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);249 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
251 dyQueue_.FreeTensor(dyLocal);250 dyQueue_.FreeTensor(dyLocal);
@@ -287,10 +286,10 @@ public:
287 xLocal = xQueue_.template DeQue<X_TYPE>();286 xLocal = xQueue_.template DeQue<X_TYPE>();
288 rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();287 rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
289 288 
290- __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();289+ __ubuf__ DY_TYPE* dyAddr = (__ubuf__ DY_TYPE*)dyLocal[0].GetPhyAddr();
291- __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();290+ __ubuf__ X_TYPE* xAddr = (__ubuf__ X_TYPE*)xLocal[0].GetPhyAddr();
292- __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();291+ __ubuf__ RSTD_TYPE* rstdAddr = (__ubuf__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
293- __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();292+ __ubuf__ float* dgammaOutAddr = (__ubuf__ float*)dgammaOutLocal[0].GetPhyAddr();
294 293 
295 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);294 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
296 dyQueue_.FreeTensor(dyLocal);295 dyQueue_.FreeTensor(dyLocal);
@@ -322,10 +321,10 @@ public:
322 xLocal = xQueue_.template DeQue<X_TYPE>();321 xLocal = xQueue_.template DeQue<X_TYPE>();
323 rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();322 rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
324 323 
325- __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();324+ __ubuf__ DY_TYPE* dyAddr = (__ubuf__ DY_TYPE*)dyLocal[0].GetPhyAddr();
326- __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();325+ __ubuf__ X_TYPE* xAddr = (__ubuf__ X_TYPE*)xLocal[0].GetPhyAddr();
327- __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();326+ __ubuf__ RSTD_TYPE* rstdAddr = (__ubuf__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
328- __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();327+ __ubuf__ float* dgammaOutAddr = (__ubuf__ float*)dgammaOutLocal[0].GetPhyAddr();
329 328 
330 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);329 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
331 dyQueue_.FreeTensor(dyLocal);330 dyQueue_.FreeTensor(dyLocal);
@@ -347,10 +346,10 @@ public:
347 xLocal1 = xQueue_.template DeQue<X_TYPE>();346 xLocal1 = xQueue_.template DeQue<X_TYPE>();
348 rstdLocal1 = rstdQueue_.template DeQue<RSTD_TYPE>();347 rstdLocal1 = rstdQueue_.template DeQue<RSTD_TYPE>();
349 348 
350- __local_mem__ DY_TYPE* dyAddr1 = (__local_mem__ DY_TYPE*)dyLocal1[0].GetPhyAddr();349+ __ubuf__ DY_TYPE* dyAddr1 = (__ubuf__ DY_TYPE*)dyLocal1[0].GetPhyAddr();
351- __local_mem__ X_TYPE* xAddr1 = (__local_mem__ X_TYPE*)xLocal1[0].GetPhyAddr();350+ __ubuf__ X_TYPE* xAddr1 = (__ubuf__ X_TYPE*)xLocal1[0].GetPhyAddr();
352- __local_mem__ RSTD_TYPE* rstdAddr1 = (__local_mem__ RSTD_TYPE*)rstdLocal1[0].GetPhyAddr();351+ __ubuf__ RSTD_TYPE* rstdAddr1 = (__ubuf__ RSTD_TYPE*)rstdLocal1[0].GetPhyAddr();
353- __local_mem__ float* dgammaOutAddr1 = (__local_mem__ float*)dgammaOutLocal1[0].GetPhyAddr();352+ __ubuf__ float* dgammaOutAddr1 = (__ubuf__ float*)dgammaOutLocal1[0].GetPhyAddr();
354 353 
355 VFCalcPreDgamma(dyAddr1, xAddr1, rstdAddr1, dgammaOutAddr1, currentCols, rowsPerUB_);354 VFCalcPreDgamma(dyAddr1, xAddr1, rstdAddr1, dgammaOutAddr1, currentCols, rowsPerUB_);
356 dyQueue_.FreeTensor(dyLocal1);355 dyQueue_.FreeTensor(dyLocal1);
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_regbase_dgamma_big_m.h+25-25
@@ -211,18 +211,18 @@ private:
211 uint32_t outerStride = td_->dgammaNfactorBlockAligned;211 uint32_t outerStride = td_->dgammaNfactorBlockAligned;
212 uint32_t innerStride = VL_FP32;212 uint32_t innerStride = VL_FP32;
213 213 
214- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();214+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
215- __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr();215+ __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr();
216- __local_mem__ T* dy = (__local_mem__ T*)dyTensor.GetPhyAddr();216+ __ubuf__ T* dy = (__ubuf__ T*)dyTensor.GetPhyAddr();
217- __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr();217+ __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr();
218 218 
219 __VEC_SCOPE__219 __VEC_SCOPE__
220 {220 {
221 for (uint16_t i = 0; i < outerLoopTimes; ++i) {221 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
222 uint32_t count = static_cast<uint32_t>(colSize);222 uint32_t count = static_cast<uint32_t>(colSize);
223 AscendC::MicroAPI::RegTensor<float> rstdReg;223 AscendC::MicroAPI::RegTensor<float> rstdReg;
224- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(224+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(
225- rstdReg, (__local_mem__ float*)rstd + static_cast<uint32_t>(i));225+ rstdReg, (__ubuf__ float*)rstd + static_cast<uint32_t>(i));
226 226 
227 AscendC::MicroAPI::RegTensor<float> xReg;227 AscendC::MicroAPI::RegTensor<float> xReg;
228 AscendC::MicroAPI::RegTensor<float> dyReg;228 AscendC::MicroAPI::RegTensor<float> dyReg;
@@ -234,7 +234,7 @@ private:
234 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask);234 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask);
235 LoadOneTensor<T>(dyReg, dy, pMask, offset);235 LoadOneTensor<T>(dyReg, dy, pMask, offset);
236 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);236 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);
237- DataCopy((__local_mem__ float*)dst + offset, xReg, pMask);237+ StoreAlign((__ubuf__ float*)dst + offset, xReg, pMask);
238 }238 }
239 }239 }
240 }240 }
@@ -249,18 +249,18 @@ private:
249 uint32_t outerStride = td_->dgammaNfactorBlockAligned;249 uint32_t outerStride = td_->dgammaNfactorBlockAligned;
250 uint32_t innerStride = VL_FP32;250 uint32_t innerStride = VL_FP32;
251 251 
252- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();252+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
253- __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr();253+ __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr();
254- __local_mem__ T* dy = (__local_mem__ T*)dyTensor.GetPhyAddr();254+ __ubuf__ T* dy = (__ubuf__ T*)dyTensor.GetPhyAddr();
255- __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr();255+ __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr();
256 256 
257 __VEC_SCOPE__257 __VEC_SCOPE__
258 {258 {
259 for (uint16_t i = 0; i < outerLoopTimes; ++i) {259 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
260 AscendC::MicroAPI::RegTensor<float> rstdReg;260 AscendC::MicroAPI::RegTensor<float> rstdReg;
261 uint32_t count = static_cast<uint32_t>(colSize);261 uint32_t count = static_cast<uint32_t>(colSize);
262- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(262+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(
263- rstdReg, (__local_mem__ float*)rstd + static_cast<uint32_t>(i));263+ rstdReg, (__ubuf__ float*)rstd + static_cast<uint32_t>(i));
264 264 
265 AscendC::MicroAPI::RegTensor<float> xReg;265 AscendC::MicroAPI::RegTensor<float> xReg;
266 AscendC::MicroAPI::RegTensor<float> dyReg;266 AscendC::MicroAPI::RegTensor<float> dyReg;
@@ -274,22 +274,22 @@ private:
274 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);274 Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);
275 LoadOneTensor<float>(dyReg, dst, pMask, offset);275 LoadOneTensor<float>(dyReg, dst, pMask, offset);
276 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);276 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);
277- DataCopy((__local_mem__ float*)dst + offset, xReg, pMask);277+ StoreAlign((__ubuf__ float*)dst + offset, xReg, pMask);
278 }278 }
279 }279 }
280 }280 }
281 }281 }
282 282 
283 template <typename T1>283 template <typename T1>
284- __aicore__ inline void LoadOneTensor(MicroAPI::RegTensor<float>& dst, const __local_mem__ void* input,284+ __aicore__ inline void LoadOneTensor(MicroAPI::RegTensor<float>& dst, const __ubuf__ void* input,
285 MicroAPI::MaskReg& preg, uint32_t offset)285 MicroAPI::MaskReg& preg, uint32_t offset)
286 {286 {
287 if constexpr (!IsSameType<T1, float>::value) {287 if constexpr (!IsSameType<T1, float>::value) {
288 MicroAPI::RegTensor<T1> xFp16;288 MicroAPI::RegTensor<T1> xFp16;
289- DataCopy<T1, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, (__local_mem__ T1*)(input) + offset);289+ LoadAlign<T1, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, (__ubuf__ T1*)(input) + offset);
290 Cast<float, T1, castTraitB162B32>(dst, xFp16, preg);290 Cast<float, T1, castTraitB162B32>(dst, xFp16, preg);
291 } else {291 } else {
292- DataCopy(dst, (__local_mem__ float*)(input) + offset);292+ LoadAlign(dst, (__ubuf__ float*)(input) + offset);
293 }293 }
294 }294 }
295 295 
@@ -350,21 +350,21 @@ private:
350 uint32_t innerLoopStride = stride;350 uint32_t innerLoopStride = stride;
351 __VEC_SCOPE__351 __VEC_SCOPE__
352 {352 {
353- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();353+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
354- __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride;354+ __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride;
355- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();355+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
356 uint32_t sreg = static_cast<uint32_t>(count);356 uint32_t sreg = static_cast<uint32_t>(count);
357 AscendC::MicroAPI::RegTensor<float> aReg, bReg;357 AscendC::MicroAPI::RegTensor<float> aReg, bReg;
358 AscendC::MicroAPI::MaskReg pMask;358 AscendC::MicroAPI::MaskReg pMask;
359 for (uint16_t i = 0; i < outerLoopTimes; ++i) {359 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
360 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);360 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
361- DataCopy(aReg, (__local_mem__ float*)src + static_cast<uint32_t>(i * outerLoopStride));361+ LoadAlign(aReg, (__ubuf__ float*)src + static_cast<uint32_t>(i * outerLoopStride));
362 for (uint16_t j = 0; j < innerLoopTimes; ++j) {362 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
363- DataCopy(bReg, (__local_mem__ float*)dst +363+ LoadAlign(bReg,
364- static_cast<uint32_t>(i * outerLoopStride + j * innerLoopStride));364+ (__ubuf__ float*)dst + static_cast<uint32_t>(i * outerLoopStride + j * innerLoopStride));
365 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);365 Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
366 }366 }
367- DataCopy((__local_mem__ float*)cah + static_cast<uint32_t>(i * outerLoopStride), aReg, pMask);367+ StoreAlign((__ubuf__ float*)cah + static_cast<uint32_t>(i * outerLoopStride), aReg, pMask);
368 }368 }
369 }369 }
370 }370 }
@@ -488,4 +488,4 @@ private:
488}; // RmsNormGradDgammaBigM488}; // RmsNormGradDgammaBigM
489 489 
490} // namespace RmsNormGrad490} // namespace RmsNormGrad
491-#endif // RMS_NORM_GRAD_DGAMMA_BIG_M_491+#endif // RMS_NORM_GRAD_DGAMMA_BIG_M_
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_regbase_dx_full_load.h+17-17
@@ -100,17 +100,17 @@ public:
100 constexpr uint32_t oneRepeat = V_LENGTH;100 constexpr uint32_t oneRepeat = V_LENGTH;
101 int64_t cols = colsAlignBlock_;101 int64_t cols = colsAlignBlock_;
102 uint16_t repeatCount = DivCeil(cols_, oneRepeat);102 uint16_t repeatCount = DivCeil(cols_, oneRepeat);
103- __local_mem__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();103+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
104- __local_mem__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();104+ __ubuf__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();
105- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();105+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
106- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();106+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
107- __local_mem__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();107+ __ubuf__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();
108 __VEC_SCOPE__108 __VEC_SCOPE__
109 {109 {
110 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;110 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;
111 for (uint16_t r = 0; r < loopRow; r++) {111 for (uint16_t r = 0; r < loopRow; r++) {
112 MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();112 MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();
113- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));113+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));
114 for (uint16_t i = 0; i < repeatCount; i++) {114 for (uint16_t i = 0; i < repeatCount; i++) {
115 LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);115 LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
116 LoadAndCast(dyReg, dyAddr, maskReg, r * cols + i * oneRepeat);116 LoadAndCast(dyReg, dyAddr, maskReg, r * cols + i * oneRepeat);
@@ -118,7 +118,7 @@ public:
118 LoadAndCast(xReg, xAddr, maskReg, r * cols + i * oneRepeat);118 LoadAndCast(xReg, xAddr, maskReg, r * cols + i * oneRepeat);
119 Mul(mulReg0, xReg, rstdReg, maskReg);119 Mul(mulReg0, xReg, rstdReg, maskReg);
120 Mul(mulReg3, mulReg2, mulReg0, maskReg);120 Mul(mulReg3, mulReg2, mulReg0, maskReg);
121- DataCopy(reduceAddr + static_cast<uint32_t>(r * colsAlign2VL_ + i * oneRepeat), mulReg3, maskReg);121+ StoreAlign(reduceAddr + static_cast<uint32_t>(r * colsAlign2VL_ + i * oneRepeat), mulReg3, maskReg);
122 }122 }
123 }123 }
124 }124 }
@@ -126,8 +126,8 @@ public:
126 MultiReduceSum(tmpSumLocal, reduceLocal, calcRowNumSub);126 MultiReduceSum(tmpSumLocal, reduceLocal, calcRowNumSub);
127 127 
128 LocalTensor<float> dxLocal = outQueueDx_.AllocTensor<float>();128 LocalTensor<float> dxLocal = outQueueDx_.AllocTensor<float>();
129- __local_mem__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();129+ __ubuf__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();
130- __local_mem__ T_X* dxAddr = (__ubuf__ T_X*)dxLocal.GetPhyAddr();130+ __ubuf__ T_X* dxAddr = (__ubuf__ T_X*)dxLocal.GetPhyAddr();
131 __VEC_SCOPE__131 __VEC_SCOPE__
132 {132 {
133 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;133 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;
@@ -135,8 +135,8 @@ public:
135 uint32_t sreg = cols_;135 uint32_t sreg = cols_;
136 int64_t cols = colsAlignBlock_;136 int64_t cols = colsAlignBlock_;
137 MaskReg maskReg = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();137 MaskReg maskReg = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
138- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));138+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));
139- DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr + static_cast<uint32_t>(r));139+ LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr + static_cast<uint32_t>(r));
140 Muls(meanReg, meanReg, avgFactor1_, maskReg);140 Muls(meanReg, meanReg, avgFactor1_, maskReg);
141 for (uint16_t i = 0; i < repeatCount; i++) {141 for (uint16_t i = 0; i < repeatCount; i++) {
142 maskReg = UpdateMask<float>(sreg);142 maskReg = UpdateMask<float>(sreg);
@@ -151,11 +151,11 @@ public:
151 Sub(subReg, mulReg2, mulReg4, maskReg);151 Sub(subReg, mulReg2, mulReg4, maskReg);
152 Mul(dxReg, subReg, rstdReg, maskReg);152 Mul(dxReg, subReg, rstdReg, maskReg);
153 if constexpr (IsSameType<T_X, float>::value) {153 if constexpr (IsSameType<T_X, float>::value) {
154- DataCopy(dxAddr + static_cast<uint32_t>(r * cols + i * oneRepeat), dxReg, maskReg);154+ StoreAlign(dxAddr + static_cast<uint32_t>(r * cols + i * oneRepeat), dxReg, maskReg);
155 } else {155 } else {
156 RegTensor<T_X> dxRegB16;156 RegTensor<T_X> dxRegB16;
157 Cast<T_X, float, castTraitB322B16>(dxRegB16, dxReg, maskReg);157 Cast<T_X, float, castTraitB322B16>(dxRegB16, dxReg, maskReg);
158- DataCopy<T_X, StoreDist::DIST_PACK_B32>(158+ StoreAlign<T_X, StoreDist::DIST_PACK_B32>(
159 dxAddr + static_cast<uint32_t>(r * cols + i * oneRepeat), dxRegB16, maskReg);159 dxAddr + static_cast<uint32_t>(r * cols + i * oneRepeat), dxRegB16, maskReg);
160 }160 }
161 }161 }
@@ -214,7 +214,7 @@ public:
214 214 
215 __aicore__ inline void MultiReduceSum(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal, int64_t rows)215 __aicore__ inline void MultiReduceSum(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal, int64_t rows)
216 {216 {
217- __local_mem__ float* srcAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();217+ __ubuf__ float* srcAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();
218 uint32_t colsTail = colsAlign2VL_ - colsAlignBlock_;218 uint32_t colsTail = colsAlign2VL_ - colsAlignBlock_;
219 constexpr uint32_t oneRepeat = V_LENGTH;219 constexpr uint32_t oneRepeat = V_LENGTH;
220 uint16_t repeatCount = DivCeil(colsTail, oneRepeat);220 uint16_t repeatCount = DivCeil(colsTail, oneRepeat);
@@ -229,8 +229,8 @@ public:
229 uint32_t sreg = colsTail;229 uint32_t sreg = colsTail;
230 for (uint16_t i = 0; i < repeatCount; i++) {230 for (uint16_t i = 0; i < repeatCount; i++) {
231 maskReg = UpdateMask<float>(sreg);231 maskReg = UpdateMask<float>(sreg);
232- DataCopy(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + cols + i * oneRepeat), srcReg,232+ StoreAlign(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + cols + i * oneRepeat), srcReg,
233- maskReg);233+ maskReg);
234 }234 }
235 }235 }
236 }236 }
@@ -281,4 +281,4 @@ private:
281 float avgFactor1_;281 float avgFactor1_;
282};282};
283} // namespace RmsNormGrad283} // namespace RmsNormGrad
284-#endif // RMS_NORM_GRAD_REGBASE_DX_FULL_LOAD_H284+#endif // RMS_NORM_GRAD_REGBASE_DX_FULL_LOAD_H
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_regbase_dx_split_d.h+21-21
@@ -196,16 +196,16 @@ public:
196 uint32_t sreg = count;196 uint32_t sreg = count;
197 constexpr uint32_t oneRepeat = V_LENGTH;197 constexpr uint32_t oneRepeat = V_LENGTH;
198 uint16_t repeatCount = DivCeil(count, oneRepeat);198 uint16_t repeatCount = DivCeil(count, oneRepeat);
199- __local_mem__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();199+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
200- __local_mem__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();200+ __ubuf__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();
201- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();201+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
202- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();202+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
203- __local_mem__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();203+ __ubuf__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();
204 __VEC_SCOPE__204 __VEC_SCOPE__
205 {205 {
206 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;206 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;
207 MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();207 MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();
208- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);208+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);
209 for (uint16_t i = 0; i < repeatCount; i++) {209 for (uint16_t i = 0; i < repeatCount; i++) {
210 LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);210 LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
211 LoadAndCast(dyReg, dyAddr, maskReg, i * oneRepeat);211 LoadAndCast(dyReg, dyAddr, maskReg, i * oneRepeat);
@@ -214,10 +214,10 @@ public:
214 Mul(mulReg0, xReg, rstdReg, maskReg);214 Mul(mulReg0, xReg, rstdReg, maskReg);
215 Mul(mulReg3, mulReg2, mulReg0, maskReg);215 Mul(mulReg3, mulReg2, mulReg0, maskReg);
216 if constexpr (IsBody) {216 if constexpr (IsBody) {
217- DataCopy(reduceAddr + static_cast<uint32_t>(i * oneRepeat), mulReg3, maskReg);217+ StoreAlign(reduceAddr + static_cast<uint32_t>(i * oneRepeat), mulReg3, maskReg);
218 } else {218 } else {
219- DataCopy(reduceAddr + static_cast<uint32_t>(ubFactorD_ + i * oneRepeat), mulReg3,219+ StoreAlign(reduceAddr + static_cast<uint32_t>(ubFactorD_ + i * oneRepeat), mulReg3,
220- maskReg); // 注意补零220+ maskReg); // 注意补零
221 }221 }
222 }222 }
223 }223 }
@@ -267,18 +267,18 @@ public:
267 uint32_t sreg = count;267 uint32_t sreg = count;
268 constexpr uint32_t oneRepeat = V_LENGTH;268 constexpr uint32_t oneRepeat = V_LENGTH;
269 uint16_t repeatCount = DivCeil(count, oneRepeat); // 可能会报错269 uint16_t repeatCount = DivCeil(count, oneRepeat); // 可能会报错
270- __local_mem__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();270+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
271- __local_mem__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();271+ __ubuf__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();
272- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();272+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
273- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();273+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
274- __local_mem__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();274+ __ubuf__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();
275- __local_mem__ T_X* dxAddr = (__ubuf__ T_X*)dxLocal.GetPhyAddr();275+ __ubuf__ T_X* dxAddr = (__ubuf__ T_X*)dxLocal.GetPhyAddr();
276 __VEC_SCOPE__276 __VEC_SCOPE__
277 {277 {
278 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;278 RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;
279 MaskReg maskReg;279 MaskReg maskReg;
280- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);280+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);
281- DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr);281+ LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr);
282 for (uint16_t i = 0; i < repeatCount; i++) {282 for (uint16_t i = 0; i < repeatCount; i++) {
283 maskReg = UpdateMask<float>(sreg);283 maskReg = UpdateMask<float>(sreg);
284 LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);284 LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
@@ -290,12 +290,12 @@ public:
290 Sub(subReg, mulReg2, mulReg4, maskReg);290 Sub(subReg, mulReg2, mulReg4, maskReg);
291 Mul(dxReg, subReg, rstdReg, maskReg);291 Mul(dxReg, subReg, rstdReg, maskReg);
292 if constexpr (IsSameType<T_X, float>::value) {292 if constexpr (IsSameType<T_X, float>::value) {
293- DataCopy(dxAddr + static_cast<uint32_t>(i * oneRepeat), dxReg, maskReg);293+ StoreAlign(dxAddr + static_cast<uint32_t>(i * oneRepeat), dxReg, maskReg);
294 } else {294 } else {
295 RegTensor<T_X> dxRegB16;295 RegTensor<T_X> dxRegB16;
296 Cast<T_X, float, castTraitB322B16>(dxRegB16, dxReg, maskReg);296 Cast<T_X, float, castTraitB322B16>(dxRegB16, dxReg, maskReg);
297- DataCopy<T_X, StoreDist::DIST_PACK_B32>(dxAddr + static_cast<uint32_t>(i * oneRepeat), dxRegB16,297+ StoreAlign<T_X, StoreDist::DIST_PACK_B32>(dxAddr + static_cast<uint32_t>(i * oneRepeat), dxRegB16,
298- maskReg);298+ maskReg);
299 }299 }
300 }300 }
301 }301 }
@@ -349,4 +349,4 @@ private:
349 float avgFactor1_;349 float avgFactor1_;
350};350};
351} // namespace RmsNormGrad351} // namespace RmsNormGrad
352-#endif // RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H352+#endif // RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H
Mnorm/rms_norm_grad_quant/op_kernel/arch35/rms_norm_grad_quant_common.h+411-413
@@ -1,413 +1,411 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/*!11+/*!
12- * \file rms_norm_grad_quant_common.h12+ * \file rms_norm_grad_quant_common.h
13- * \brief RmsNormGradQuant common13+ * \brief RmsNormGradQuant common
14- */14+ */
15-#ifndef OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_QUANT_COMMON_H15+#ifndef OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_QUANT_COMMON_H
16-#define OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_QUANT_COMMON_H16+#define OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_QUANT_COMMON_H
17-#define FLOAT_OVERFLOW_MODE_CTRL 6017+#define FLOAT_OVERFLOW_MODE_CTRL 60
18- 18+ 
19-#include "kernel_operator.h"19+#include "kernel_operator.h"
20-#include "../inc/platform.h"20+#include "../inc/platform.h"
21-#include "rms_norm_grad_quant_tiling_data.h"21+#include "rms_norm_grad_quant_tiling_data.h"
22- 22+ 
23-namespace RmsNormGradQuant {23+namespace RmsNormGradQuant {
24-using namespace AscendC;24+using namespace AscendC;
25-using namespace AscendC::MicroAPI;25+using namespace AscendC::MicroAPI;
26-using AscendC::MicroAPI::CreateMask;26+using AscendC::MicroAPI::CreateMask;
27-using AscendC::MicroAPI::LoadDist;27+using AscendC::MicroAPI::LoadDist;
28-using AscendC::MicroAPI::LocalMemBar;28+using AscendC::MicroAPI::LocalMemBar;
29-using AscendC::MicroAPI::MaskPattern;29+using AscendC::MicroAPI::MaskPattern;
30-using AscendC::MicroAPI::MaskReg;30+using AscendC::MicroAPI::MaskReg;
31-using AscendC::MicroAPI::MemType;31+using AscendC::MicroAPI::MemType;
32-using AscendC::MicroAPI::RegTensor;32+using AscendC::MicroAPI::RegTensor;
33-using AscendC::MicroAPI::StoreDist;33+using AscendC::MicroAPI::StoreDist;
34-using AscendC::MicroAPI::UpdateMask;34+using AscendC::MicroAPI::UpdateMask;
35- 35+ 
36-namespace RmsNormGradRegbase {36+namespace RmsNormGradRegbase {
37-__aicore__ inline constexpr uint32_t GetVRegSize()37+__aicore__ inline constexpr uint32_t GetVRegSize()
38-{38+{
39-#if __CCE_AICORE__ == 31039+#if __CCE_AICORE__ == 310
40- return AscendC::VECTOR_REG_WIDTH;40+ return AscendC::VECTOR_REG_WIDTH;
41-#else41+#else
42- return 256U;42+ return 256U;
43-#endif43+#endif
44-}44+}
45-} // namespace RmsNormGradRegbase45+} // namespace RmsNormGradRegbase
46- 46+ 
47-constexpr uint32_t ONCE_VECTOR_SIZE = 256;47+constexpr uint32_t ONCE_VECTOR_SIZE = 256;
48-constexpr uint32_t V_LENGTH = RmsNormGradRegbase::GetVRegSize() / sizeof(float);48+constexpr uint32_t V_LENGTH = RmsNormGradRegbase::GetVRegSize() / sizeof(float);
49-constexpr uint32_t FLOAT_NUM_BLOCK = 8;49+constexpr uint32_t FLOAT_NUM_BLOCK = 8;
50-constexpr uint32_t HALF_NUM_BLOCK = 16;50+constexpr uint32_t HALF_NUM_BLOCK = 16;
51-constexpr uint32_t HIFP8_NUM_BLOCK = 32;51+constexpr uint32_t HIFP8_NUM_BLOCK = 32;
52-constexpr uint32_t FLOAT_NUM_2VL = 128;52+constexpr uint32_t FLOAT_NUM_2VL = 128;
53-constexpr uint32_t DB_NUM = 2;53+constexpr uint32_t DB_NUM = 2;
54-constexpr uint32_t DEPTH_TWO = 2;54+constexpr uint32_t DEPTH_TWO = 2;
55-constexpr uint32_t NUM_ONE = 1;55+constexpr uint32_t NUM_ONE = 1;
56-constexpr uint32_t NUM_TWO = 2;56+constexpr uint32_t NUM_TWO = 2;
57-constexpr uint32_t UB_FACTOR_DX_FULL_LOAD = 6144;57+constexpr uint32_t UB_FACTOR_DX_FULL_LOAD = 6144;
58-constexpr uint32_t UB_FACTOR_DX_SPLIT_D = 4096;58+constexpr uint32_t UB_FACTOR_DX_SPLIT_D = 4096;
59-constexpr uint32_t BLOCKSIZE = platform::GetUbBlockSize();59+constexpr uint32_t BLOCKSIZE = platform::GetUbBlockSize();
60-constexpr uint32_t BLOCKSIZEB32 = platform::GetUbBlockSize() / sizeof(float);60+constexpr uint32_t BLOCKSIZEB32 = platform::GetUbBlockSize() / sizeof(float);
61- 61+ 
62-constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {62+constexpr AscendC::MicroAPI::CastTrait castTraitB162B32 = {
63- AscendC::MicroAPI::RegLayout::ZERO,63+ AscendC::MicroAPI::RegLayout::ZERO,
64- AscendC::MicroAPI::SatMode::UNKNOWN,64+ AscendC::MicroAPI::SatMode::UNKNOWN,
65- AscendC::MicroAPI::MaskMergeMode::ZEROING,65+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
66- AscendC::RoundMode::UNKNOWN,66+ AscendC::RoundMode::UNKNOWN,
67-};67+};
68- 68+ 
69-constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = {69+constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = {
70- AscendC::MicroAPI::RegLayout::ZERO,70+ AscendC::MicroAPI::RegLayout::ZERO,
71- AscendC::MicroAPI::SatMode::NO_SAT,71+ AscendC::MicroAPI::SatMode::NO_SAT,
72- AscendC::MicroAPI::MaskMergeMode::ZEROING,72+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
73- AscendC::RoundMode::CAST_RINT,73+ AscendC::RoundMode::CAST_RINT,
74-};74+};
75- 75+ 
76-constexpr AscendC::MicroAPI::CastTrait castTraitFp322Hifp8 = {76+constexpr AscendC::MicroAPI::CastTrait castTraitFp322Hifp8 = {
77- AscendC::MicroAPI::RegLayout::ZERO,77+ AscendC::MicroAPI::RegLayout::ZERO,
78- AscendC::MicroAPI::SatMode::SAT,78+ AscendC::MicroAPI::SatMode::SAT,
79- AscendC::MicroAPI::MaskMergeMode::ZEROING,79+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
80- RoundMode::CAST_ROUND,80+ RoundMode::CAST_ROUND,
81-};81+};
82- 82+ 
83-constexpr AscendC::MicroAPI::CastTrait castTraitFp322Int32 = {83+constexpr AscendC::MicroAPI::CastTrait castTraitFp322Int32 = {
84- AscendC::MicroAPI::RegLayout::UNKNOWN,84+ AscendC::MicroAPI::RegLayout::UNKNOWN,
85- AscendC::MicroAPI::SatMode::NO_SAT,85+ AscendC::MicroAPI::SatMode::NO_SAT,
86- AscendC::MicroAPI::MaskMergeMode::ZEROING,86+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
87- AscendC::RoundMode::CAST_RINT,87+ AscendC::RoundMode::CAST_RINT,
88-};88+};
89- 89+ 
90-constexpr AscendC::MicroAPI::CastTrait castTraitFp322Fp16 = {90+constexpr AscendC::MicroAPI::CastTrait castTraitFp322Fp16 = {
91- AscendC::MicroAPI::RegLayout::ZERO,91+ AscendC::MicroAPI::RegLayout::ZERO,
92- AscendC::MicroAPI::SatMode::NO_SAT,92+ AscendC::MicroAPI::SatMode::NO_SAT,
93- AscendC::MicroAPI::MaskMergeMode::ZEROING,93+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
94- AscendC::RoundMode::CAST_RINT,94+ AscendC::RoundMode::CAST_RINT,
95-};95+};
96- 96+ 
97-constexpr AscendC::MicroAPI::CastTrait castTraitFp162Int8 = {97+constexpr AscendC::MicroAPI::CastTrait castTraitFp162Int8 = {
98- AscendC::MicroAPI::RegLayout::ZERO,98+ AscendC::MicroAPI::RegLayout::ZERO,
99- AscendC::MicroAPI::SatMode::NO_SAT,99+ AscendC::MicroAPI::SatMode::NO_SAT,
100- AscendC::MicroAPI::MaskMergeMode::ZEROING,100+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
101- AscendC::RoundMode::CAST_TRUNC,101+ AscendC::RoundMode::CAST_TRUNC,
102-};102+};
103- 103+ 
104-constexpr AscendC::MicroAPI::CastTrait castTraitInt322Fp32 = {104+constexpr AscendC::MicroAPI::CastTrait castTraitInt322Fp32 = {
105- AscendC::MicroAPI::RegLayout::UNKNOWN,105+ AscendC::MicroAPI::RegLayout::UNKNOWN,
106- AscendC::MicroAPI::SatMode::NO_SAT,106+ AscendC::MicroAPI::SatMode::NO_SAT,
107- AscendC::MicroAPI::MaskMergeMode::ZEROING,107+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
108- AscendC::RoundMode::CAST_RINT,108+ AscendC::RoundMode::CAST_RINT,
109-};109+};
110- 110+ 
111-template <typename T>111+template <typename T>
112-__aicore__ inline T Min(T left, T right)112+__aicore__ inline T Min(T left, T right)
113-{113+{
114- return (left < right ? left : right);114+ return (left < right ? left : right);
115-}115+}
116- 116+ 
117-__aicore__ inline int32_t findPowerTwo(int32_t n)117+__aicore__ inline int32_t findPowerTwo(int32_t n)
118-{118+{
119- // find max power of 2 no more than n (32 bit)119+ // find max power of 2 no more than n (32 bit)
120- n |= n >> 1; // 将最高位的1向右扩展1位， 11XX..X120+ n |= n >> 1; // 将最高位的1向右扩展1位， 11XX..X
121- n |= n >> 2; // 将最高位的1向右扩展2位， 1111X..X121+ n |= n >> 2; // 将最高位的1向右扩展2位， 1111X..X
122- n |= n >> 4; // 将最高位的1向右扩展4位122+ n |= n >> 4; // 将最高位的1向右扩展4位
123- n |= n >> 8; // 将最高位的1向右扩展8位123+ n |= n >> 8; // 将最高位的1向右扩展8位
124- n |= n >> 16; // 将最高位的1向右扩展16位124+ n |= n >> 16; // 将最高位的1向右扩展16位
125- return (n + 1) >> 1;125+ return (n + 1) >> 1;
126-}126+}
127- 127+ 
128-template <typename T>128+template <typename T>
129-__aicore__ inline void LoadAndCast(RegTensor<float>& dstReg, __local_mem__ T* srcAddr, MaskReg& maskReg,129+__aicore__ inline void LoadAndCast(RegTensor<float>& dstReg, __ubuf__ T* srcAddr, MaskReg& maskReg, uint32_t srcOffset)
130- uint32_t srcOffset)130+{
131-{131+ if constexpr (IsSameType<T, float>::value) {
132- if constexpr (IsSameType<T, float>::value) {132+ LoadAlign(dstReg, srcAddr + srcOffset);
133- DataCopy(dstReg, srcAddr + srcOffset);133+ } else {
134- } else {134+ RegTensor<T> dstRegB16;
135- RegTensor<T> dstRegB16;135+ // LoadAlign<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + ubFactorD_ + srcOffset); // 后续不做偏移
136- // DataCopy<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + ubFactorD_ + srcOffset); // 后续不做偏移136+ LoadAlign<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + srcOffset); // 后续不做偏移
137- DataCopy<T, LoadDist::DIST_UNPACK_B16>(dstRegB16, srcAddr + srcOffset); // 后续不做偏移137+ Cast<float, T, castTraitB162B32>(dstReg, dstRegB16, maskReg);
138- Cast<float, T, castTraitB162B32>(dstReg, dstRegB16, maskReg);138+ }
139- }139+}
140-}140+ 
141- 141+/*!
142-/*!142+ * Compute ReduceSum mean
143- * Compute ReduceSum mean143+ *
144- *144+ * @param dstLocal dst levelTensor
145- * @param dstLocal dst levelTensor145+ * @param srcLocal src LevelTenso
146- * @param srcLocal src LevelTenso146+ * @param offset dst offset
147- * @param offset dst offset147+ * @param count src level size
148- * @param count src level size148+ * @return
149- * @return149+ */
150- */150+__aicore__ inline void LevelMerge(LocalTensor<float>& dstLocal, LocalTensor<float> srcLocal, uint64_t offset,
151-__aicore__ inline void LevelMerge(LocalTensor<float>& dstLocal, LocalTensor<float> srcLocal, uint64_t offset,151+ uint32_t count)
152- uint32_t count)152+{
153-{153+ uint64_t calCount = count / 4;
154- uint64_t calCount = count / 4;154+ uint32_t sreg = (uint32_t)(calCount);
155- uint32_t sreg = (uint32_t)(calCount);155+ uint16_t repeatTimes = CeilDivision(calCount, V_LENGTH);
156- uint16_t repeatTimes = CeilDivision(calCount, V_LENGTH);156+ uint32_t meanTile = repeatTimes;
157- uint32_t meanTile = repeatTimes;157+ 
158- 158+ __ubuf__ float* src1Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 0 * calCount;
159- __local_mem__ float* src1Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 0 * calCount;159+ __ubuf__ float* src2Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 1 * calCount;
160- __local_mem__ float* src2Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 1 * calCount;160+ __ubuf__ float* src3Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 2 * calCount;
161- __local_mem__ float* src3Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 2 * calCount;161+ __ubuf__ float* src4Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 3 * calCount;
162- __local_mem__ float* src4Addr = (__ubuf__ float*)srcLocal.GetPhyAddr() + 3 * calCount;162+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
163- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();163+ 
164- 164+ __VEC_SCOPE__
165- __VEC_SCOPE__165+ {
166- {166+ RegTensor<float> vRegA, vRegB, vRegC, vRegD, dstReg, vMean;
167- RegTensor<float> vRegA, vRegB, vRegC, vRegD, dstReg, vMean;167+ MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();
168- MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();168+ MaskReg pregLoop;
169- MaskReg pregLoop;169+ for (uint16_t i = 0; i < repeatTimes; ++i) {
170- for (uint16_t i = 0; i < repeatTimes; ++i) {170+ pregLoop = UpdateMask<float>(sreg);
171- pregLoop = UpdateMask<float>(sreg);171+ LoadAlign(vRegA, src1Addr + static_cast<uint32_t>(i * V_LENGTH));
172- DataCopy(vRegA, src1Addr + static_cast<uint32_t>(i * V_LENGTH));172+ LoadAlign(vRegB, src2Addr + static_cast<uint32_t>(i * V_LENGTH));
173- DataCopy(vRegB, src2Addr + static_cast<uint32_t>(i * V_LENGTH));173+ LoadAlign(vRegC, src3Addr + static_cast<uint32_t>(i * V_LENGTH));
174- DataCopy(vRegC, src3Addr + static_cast<uint32_t>(i * V_LENGTH));174+ LoadAlign(vRegD, src4Addr + static_cast<uint32_t>(i * V_LENGTH));
175- DataCopy(vRegD, src4Addr + static_cast<uint32_t>(i * V_LENGTH));175+ Add(vRegA, vRegA, vRegB, pregLoop);
176- Add(vRegA, vRegA, vRegB, pregLoop);176+ Add(vRegC, vRegC, vRegD, pregLoop);
177- Add(vRegC, vRegC, vRegD, pregLoop);177+ Add(dstReg, vRegA, vRegC, pregLoop);
178- Add(dstReg, vRegA, vRegC, pregLoop);178+ Reduce<ReduceType::SUM>(vMean, dstReg, pregLoop);
179- ReduceSum(vMean, dstReg, pregLoop);179+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(offset), vMean,
180- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(offset), vMean,180+ pregMerge);
181- pregMerge);181+ }
182- }182+ }
183- }183+}
184-}184+ 
185- 185+/*!
186-/*!186+ * The num of each level elements is 256, ReduceSum these elements and store to the next level.
187- * The num of each level elements is 256, ReduceSum these elements and store to the next level.187+ * @param level1Local level1Tensor
188- * @param level1Local level1Tensor188+ * @param level2Local level2Tensor
189- * @param level2Local level2Tensor189+ * @param level3Local level3Tensor
190- * @param level3Local level3Tensor190+ * @param level1 level1 elements
191- * @param level1 level1 elements191+ * @param level2 level2 elements
192- * @param level2 level2 elements192+ * @param level3 level3 elements
193- * @param level3 level3 elements193+ * @return void
194- * @return void194+ */
195- */195+__aicore__ inline void ComputeMultiLevelReduce(LocalTensor<float>& level1Local, LocalTensor<float>& level2Local,
196-__aicore__ inline void ComputeMultiLevelReduce(LocalTensor<float>& level1Local, LocalTensor<float>& level2Local,196+ LocalTensor<float>& level3Local, uint32_t& level1, uint32_t& level2,
197- LocalTensor<float>& level3Local, uint32_t& level1, uint32_t& level2,197+ uint32_t& level3)
198- uint32_t& level3)198+{
199-{199+ if (level1 == ONCE_VECTOR_SIZE) {
200- if (level1 == ONCE_VECTOR_SIZE) {200+ LevelMerge(level2Local, level1Local, level2, ONCE_VECTOR_SIZE);
201- LevelMerge(level2Local, level1Local, level2, ONCE_VECTOR_SIZE);201+ level1 = 0;
202- level1 = 0;202+ level2 += 1;
203- level2 += 1;203+ }
204- }204+ if (level2 == ONCE_VECTOR_SIZE) {
205- if (level2 == ONCE_VECTOR_SIZE) {205+ LevelMerge(level3Local, level2Local, level3, ONCE_VECTOR_SIZE);
206- LevelMerge(level3Local, level2Local, level3, ONCE_VECTOR_SIZE);206+ level2 = 0;
207- level2 = 0;207+ level3 += 1;
208- level3 += 1;208+ }
209- }209+}
210-}210+ 
211- 211+/*!
212-/*!212+ * compute final ReduceSum result
213- * compute final ReduceSum result213+ * @param dstLocal dst Tensor
214- * @param dstLocal dst Tensor214+ * @param offset dst offset
215- * @param offset dst offset215+ * @param level1Local level1 Tensor
216- * @param level1Local level1 Tensor216+ * @param level2Local level2 Tensor
217- * @param level2Local level2 Tensor217+ * @param level3Local level3 Tensor
218- * @param level3Local level3 Tensor218+ * @param level1 level1 elements
219- * @param level1 level1 elements219+ * @param level2 level2 elements
220- * @param level2 level2 elements220+ * @param level3 level3 elements
221- * @param level3 level3 elements221+ * @return
222- * @return222+ */
223- */223+__aicore__ inline void ComputeMultiLevelMean(LocalTensor<float>& dstLocal, uint32_t offset,
224-__aicore__ inline void ComputeMultiLevelMean(LocalTensor<float>& dstLocal, uint32_t offset,224+ LocalTensor<float>& level1Local, LocalTensor<float>& level2Local,
225- LocalTensor<float>& level1Local, LocalTensor<float>& level2Local,225+ LocalTensor<float>& level3Local, uint32_t& level1, uint32_t& level2)
226- LocalTensor<float>& level3Local, uint32_t& level1, uint32_t& level2)226+{
227-{227+ if (level1 > 0 && level1 < ONCE_VECTOR_SIZE) {
228- if (level1 > 0 && level1 < ONCE_VECTOR_SIZE) {228+ LevelMerge(dstLocal, level1Local, offset, ONCE_VECTOR_SIZE);
229- LevelMerge(dstLocal, level1Local, offset, ONCE_VECTOR_SIZE);229+ } else if (level2 > 0 && level2 < ONCE_VECTOR_SIZE) {
230- } else if (level2 > 0 && level2 < ONCE_VECTOR_SIZE) {230+ LevelMerge(dstLocal, level2Local, offset, ONCE_VECTOR_SIZE);
231- LevelMerge(dstLocal, level2Local, offset, ONCE_VECTOR_SIZE);231+ } else {
232- } else {232+ LevelMerge(dstLocal, level3Local, offset, ONCE_VECTOR_SIZE);
233- LevelMerge(dstLocal, level3Local, offset, ONCE_VECTOR_SIZE);233+ }
234- }234+}
235-}235+ 
236- 236+/*!
237-/*!237+ * ReduceSum impl by half add.
238- * ReduceSum impl by half add.238+ * @param dstLocal dst Tensor
239- * @param dstLocal dst Tensor239+ * @param srcLocal src Tensor
240- * @param srcLocal src Tensor240+ * @param workLocal temp Tensor
241- * @param workLocal temp Tensor241+ * @param offset dst offset
242- * @param offset dst offset242+ * @param count count aligned compute elements.
243- * @param count count aligned compute elements.243+ * @param powerSplit 2 ** k = powerSplit
244- * @param powerSplit 2 ** k = powerSplit244+ * @return void
245- * @return void245+ */
246- */246+__aicore__ inline void ReduceSumImpl(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal,
247-__aicore__ inline void ReduceSumImpl(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal,247+ LocalTensor<float>& workLocal, uint32_t offset, uint32_t count,
248- LocalTensor<float>& workLocal, uint32_t offset, uint32_t count,248+ uint32_t powerSplit)
249- uint32_t powerSplit)249+{
250-{250+ uint32_t remainTile = count - powerSplit;
251- uint32_t remainTile = count - powerSplit;251+ uint32_t remainSreg = remainTile;
252- uint32_t remainSreg = remainTile;252+ uint32_t remainRepeats = remainTile / (2 * V_LENGTH);
253- uint32_t remainRepeats = remainTile / (2 * V_LENGTH);253+ 
254- 254+ uint32_t masterTile = powerSplit - remainTile;
255- uint32_t masterTile = powerSplit - remainTile;255+ uint32_t masterSreg = masterTile;
256- uint32_t masterSreg = masterTile;256+ uint16_t masterRepeats = masterTile / (2 * V_LENGTH);
257- uint16_t masterRepeats = masterTile / (2 * V_LENGTH);257+ 
258- 258+ uint32_t mergeTile = powerSplit / (2 * V_LENGTH);
259- uint32_t mergeTile = powerSplit / (2 * V_LENGTH);259+ uint32_t mergeSreg = mergeTile;
260- uint32_t mergeSreg = mergeTile;260+ uint32_t mergeRepeats = mergeTile / (2 * V_LENGTH);
261- uint32_t mergeRepeats = mergeTile / (2 * V_LENGTH);261+ 
262- 262+ uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;
263- uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;263+ uint32_t meanSreg = meanTile;
264- uint32_t meanSreg = meanTile;264+ 
265- 265+ __ubuf__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();
266- __local_mem__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();266+ __ubuf__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(powerSplit);
267- __local_mem__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(powerSplit);267+ __ubuf__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(remainTile);
268- __local_mem__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + int64_t(remainTile);268+ __ubuf__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();
269- __local_mem__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();269+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
270- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();270+ 
271- 271+ __VEC_SCOPE__
272- __VEC_SCOPE__272+ {
273- {273+ RegTensor<float> mainA, mainB, tailA, tailB, vMean;
274- RegTensor<float> mainA, mainB, tailA, tailB, vMean;274+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
275- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();275+ MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();
276- MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();276+ MaskReg pregLoop;
277- MaskReg pregLoop;277+ 
278- 278+ for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {
279- for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {279+ pregLoop = UpdateMask<float>(remainSreg);
280- pregLoop = UpdateMask<float>(remainSreg);280+ LoadAlign(mainA, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
281- DataCopy(mainA, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));281+ LoadAlign(mainB, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
282- DataCopy(mainB, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));282+ LoadAlign(tailA, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
283- DataCopy(tailA, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));283+ LoadAlign(tailB, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
284- DataCopy(tailB, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));284+ 
285- 285+ Add(mainA, mainA, tailA, pregLoop);
286- Add(mainA, mainA, tailA, pregLoop);286+ Add(mainB, mainB, tailB, pregLoop);
287- Add(mainB, mainB, tailB, pregLoop);287+ Add(mainA, mainA, mainB, pregLoop);
288- Add(mainA, mainA, mainB, pregLoop);288+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
289- ReduceSum(vMean, mainA, pregLoop);289+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean, pregMerge);
290- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean, pregMerge);290+ }
291- }291+ for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
292- for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {292+ pregLoop = UpdateMask<float>(masterSreg);
293- pregLoop = UpdateMask<float>(masterSreg);293+ LoadAlign(mainA, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
294- DataCopy(mainA, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));294+ LoadAlign(mainB, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
295- DataCopy(mainB, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));295+ Add(mainA, mainA, mainB, pregLoop);
296- Add(mainA, mainA, mainB, pregLoop);296+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
297- ReduceSum(vMean, mainA, pregLoop);297+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(remainRepeats + i),
298- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(remainRepeats + i),298+ vMean, pregMerge);
299- vMean, pregMerge);299+ }
300- }300+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
301- LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();301+ for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
302- for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {302+ pregLoop = UpdateMask<float>(mergeSreg);
303- pregLoop = UpdateMask<float>(mergeSreg);303+ LoadAlign(mainA, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
304- DataCopy(mainA, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));304+ LoadAlign(mainB, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
305- DataCopy(mainB, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));305+ Add(mainA, mainA, mainB, pregLoop);
306- Add(mainA, mainA, mainB, pregLoop);306+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
307- ReduceSum(vMean, mainA, pregLoop);307+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean, pregMerge);
308- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean, pregMerge);308+ }
309- }309+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
310- LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();310+ {
311- {311+ pregLoop = UpdateMask<float>(meanSreg);
312- pregLoop = UpdateMask<float>(meanSreg);312+ LoadAlign(mainA, workAddr + 0);
313- DataCopy(mainA, workAddr + 0);313+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
314- ReduceSum(vMean, mainA, pregLoop);314+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMean, pregMerge);
315- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + offset, vMean, pregMerge);315+ }
316- }316+ }
317- }317+}
318-}318+ 
319- 319+/*!
320-/*!320+ * MultiReduceSum(AR pattern) impl by half add.
321- * MultiReduceSum(AR pattern) impl by half add.321+ * @param dstLocal dst Tensor
322- * @param dstLocal dst Tensor322+ * @param srcLocal src Tensor
323- * @param srcLocal src Tensor323+ * @param workLocal temp Tensor
324- * @param workLocal temp Tensor324+ * @param rows calc rows once
325- * @param rows calc rows once325+ * @param colsAlign2VL cols aligned 512B.
326- * @param colsAlign2VL cols aligned 512B.326+ * @param powerSplit 2 ** k = powerSplit
327- * @param powerSplit 2 ** k = powerSplit327+ * @return void
328- * @return void328+ */
329- */329+__aicore__ inline void MultiReduceSumImpl(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal,
330-__aicore__ inline void MultiReduceSumImpl(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal,330+ LocalTensor<float>& workLocal, uint32_t rows, uint32_t colsAlign2VL,
331- LocalTensor<float>& workLocal, uint32_t rows, uint32_t colsAlign2VL,331+ uint32_t powerSplit)
332- uint32_t powerSplit)332+{
333-{333+ uint32_t remainTile = colsAlign2VL - powerSplit;
334- uint32_t remainTile = colsAlign2VL - powerSplit;334+ uint32_t remainRepeats = remainTile / (2 * V_LENGTH);
335- uint32_t remainRepeats = remainTile / (2 * V_LENGTH);335+ 
336- 336+ uint32_t masterTile = powerSplit - remainTile;
337- uint32_t masterTile = powerSplit - remainTile;337+ uint16_t masterRepeats = masterTile / (2 * V_LENGTH);
338- uint16_t masterRepeats = masterTile / (2 * V_LENGTH);338+ 
339- 339+ uint32_t mergeTile = powerSplit / (2 * V_LENGTH);
340- uint32_t mergeTile = powerSplit / (2 * V_LENGTH);340+ uint32_t mergeRepeats = mergeTile / (2 * V_LENGTH);
341- uint32_t mergeRepeats = mergeTile / (2 * V_LENGTH);341+ 
342- 342+ uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;
343- uint32_t meanTile = mergeRepeats == 0 ? mergeTile : mergeRepeats;343+ 
344- 344+ __ubuf__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();
345- __local_mem__ float* workAddr = (__ubuf__ float*)workLocal.GetPhyAddr();345+ __ubuf__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();
346- __local_mem__ float* dstAddr = (__ubuf__ float*)dstLocal.GetPhyAddr();346+ 
347- 347+ __VEC_SCOPE__
348- __VEC_SCOPE__348+ {
349- {349+ RegTensor<float> mainA, mainB, tailA, tailB, vMean;
350- RegTensor<float> mainA, mainB, tailA, tailB, vMean;350+ for (uint16_t r = 0; r < (uint16_t)rows; ++r) {
351- for (uint16_t r = 0; r < (uint16_t)rows; ++r) {351+ uint32_t remainSreg = remainTile;
352- uint32_t remainSreg = remainTile;352+ uint32_t masterSreg = masterTile;
353- uint32_t masterSreg = masterTile;353+ uint32_t mergeSreg = mergeTile;
354- uint32_t mergeSreg = mergeTile;354+ uint32_t meanSreg = meanTile;
355- uint32_t meanSreg = meanTile;355+ 
356- 356+ __ubuf__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL;
357- __local_mem__ float* mainAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL;357+ __ubuf__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL + int64_t(powerSplit);
358- __local_mem__ float* tailAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL +358+ __ubuf__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL +
359- int64_t(powerSplit);359+ int64_t(remainTile);
360- __local_mem__ float* masterAddr = (__ubuf__ float*)srcLocal.GetPhyAddr() + r * colsAlign2VL +360+ 
361- int64_t(remainTile);361+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
362- 362+ MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();
363- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();363+ MaskReg pregLoop;
364- MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>();364+ 
365- MaskReg pregLoop;365+ for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {
366- 366+ pregLoop = UpdateMask<float>(remainSreg);
367- for (uint16_t i = 0; i < (uint16_t)remainRepeats; ++i) {367+ LoadAlign(mainA, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
368- pregLoop = UpdateMask<float>(remainSreg);368+ LoadAlign(mainB, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
369- DataCopy(mainA, mainAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));369+ LoadAlign(tailA, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
370- DataCopy(mainB, mainAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));370+ LoadAlign(tailB, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
371- DataCopy(tailA, tailAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));371+ 
372- DataCopy(tailB, tailAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));372+ Add(mainA, mainA, tailA, pregLoop);
373- 373+ Add(mainB, mainB, tailB, pregLoop);
374- Add(mainA, mainA, tailA, pregLoop);374+ Add(mainA, mainA, mainB, pregLoop);
375- Add(mainB, mainB, tailB, pregLoop);375+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
376- Add(mainA, mainA, mainB, pregLoop);376+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,
377- ReduceSum(vMean, mainA, pregLoop);377+ pregMerge);
378- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,378+ }
379- pregMerge);379+ for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {
380- }380+ pregLoop = UpdateMask<float>(masterSreg);
381- for (uint16_t i = 0; i < (uint16_t)masterRepeats; ++i) {381+ LoadAlign(mainA, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
382- pregLoop = UpdateMask<float>(masterSreg);382+ LoadAlign(mainB, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
383- DataCopy(mainA, masterAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));383+ Add(mainA, mainA, mainB, pregLoop);
384- DataCopy(mainB, masterAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));384+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
385- Add(mainA, mainA, mainB, pregLoop);385+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(
386- ReduceSum(vMean, mainA, pregLoop);386+ workAddr + static_cast<uint32_t>(remainRepeats + i), vMean, pregMerge);
387- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(remainRepeats + i),387+ }
388- vMean, pregMerge);388+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
389- }389+ for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {
390- LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();390+ pregLoop = UpdateMask<float>(mergeSreg);
391- for (uint16_t i = 0; i < (uint16_t)mergeRepeats; ++i) {391+ LoadAlign(mainA, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));
392- pregLoop = UpdateMask<float>(mergeSreg);392+ LoadAlign(mainB, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));
393- DataCopy(mainA, workAddr + static_cast<uint32_t>((i * 2 + 0) * V_LENGTH));393+ Add(mainA, mainA, mainB, pregLoop);
394- DataCopy(mainB, workAddr + static_cast<uint32_t>((i * 2 + 1) * V_LENGTH));394+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
395- Add(mainA, mainA, mainB, pregLoop);395+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,
396- ReduceSum(vMean, mainA, pregLoop);396+ pregMerge);
397- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(workAddr + static_cast<uint32_t>(i), vMean,397+ }
398- pregMerge);398+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
399- }399+ {
400- LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();400+ pregLoop = UpdateMask<float>(meanSreg);
401- {401+ LoadAlign(mainA, workAddr + 0);
402- pregLoop = UpdateMask<float>(meanSreg);402+ Reduce<ReduceType::SUM>(vMean, mainA, pregLoop);
403- DataCopy(mainA, workAddr + 0);403+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(r), vMean,
404- ReduceSum(vMean, mainA, pregLoop);404+ pregMerge);
405- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstAddr + static_cast<uint32_t>(r), vMean,405+ }
406- pregMerge);406+ }
407- }407+ }
408- }408+}
409- }409+ 
410-}410+} // namespace RmsNormGradQuant
411- 411+#endif // OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_QUANT_COMMON_H
412-} // namespace RmsNormGradQuant
413-#endif // OPS_BUILT_IN_TBE_IMPL_ASCENDC_RMS_NORM_GRAD_QUANT_COMMON_H
Mnorm/rms_norm_grad_quant/op_kernel/arch35/rms_norm_grad_quant_dgamma.h+515-516
@@ -1,516 +1,515 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/*!11+/*!
12- * \file rms_norm_grad_quant_regbase_dgamma.h12+ * \file rms_norm_grad_quant_regbase_dgamma.h
13- * \brief RmsNormGradQuant regbase dgamma13+ * \brief RmsNormGradQuant regbase dgamma
14- */14+ */
15- 15+ 
16-#ifndef RMS_NORM_GRAD_QUANT_DGAMMA_H16+#ifndef RMS_NORM_GRAD_QUANT_DGAMMA_H
17-#define RMS_NORM_GRAD_QUANT_DGAMMA_H17+#define RMS_NORM_GRAD_QUANT_DGAMMA_H
18- 18+ 
19-#include "kernel_tiling/kernel_tiling.h"19+#include "kernel_tiling/kernel_tiling.h"
20-#include "kernel_operator.h"20+#include "kernel_operator.h"
21-#include "rms_norm_grad_quant_dgamma_helper.h"21+#include "rms_norm_grad_quant_dgamma_helper.h"
22- 22+ 
23-namespace RmsNormGradQuant {23+namespace RmsNormGradQuant {
24-using namespace AscendC;24+using namespace AscendC;
25-template <typename DY_TYPE, typename X_TYPE, typename RSTD_TYPE, bool IS_FULL_LOAD, int32_t BUFFER_NUM = 2>25+template <typename DY_TYPE, typename X_TYPE, typename RSTD_TYPE, bool IS_FULL_LOAD, int32_t BUFFER_NUM = 2>
26-class RegbaseDgamma {26+class RegbaseDgamma {
27-public:27+public:
28- __aicore__ inline RegbaseDgamma(TPipe* pipe, const RmsNormGradQuantRegbaseTilingData* tilingData)28+ __aicore__ inline RegbaseDgamma(TPipe* pipe, const RmsNormGradQuantRegbaseTilingData* tilingData)
29- : Ppipe_(pipe), tiling_(tilingData)29+ : Ppipe_(pipe), tiling_(tilingData)
30- {}30+ {}
31- 31+ 
32- __aicore__ inline void Init(__gm__ uint8_t* dy, __gm__ uint8_t* x, __gm__ uint8_t* rstd, __gm__ uint8_t* gamma,32+ __aicore__ inline void Init(__gm__ uint8_t* dy, __gm__ uint8_t* x, __gm__ uint8_t* rstd, __gm__ uint8_t* gamma,
33- __gm__ uint8_t* dx, __gm__ uint8_t* dgamma)33+ __gm__ uint8_t* dx, __gm__ uint8_t* dgamma)
34- {34+ {
35- coreIdx_ = AscendC::GetBlockIdx();35+ coreIdx_ = AscendC::GetBlockIdx();
36- if (coreIdx_ >= tiling_->usedCoreNumDG) {36+ if (coreIdx_ >= tiling_->usedCoreNumDG) {
37- return;37+ return;
38- }38+ }
39- blockSize_ = tiling_->blockSize;39+ blockSize_ = tiling_->blockSize;
40- vlFp32_ = VECTOR_REG_WIDTH / sizeof(float);40+ vlFp32_ = VECTOR_REG_WIDTH / sizeof(float);
41- colsPerUB_ = vlFp32_;41+ colsPerUB_ = vlFp32_;
42- cols_ = tiling_->dxTilingData.cols;42+ cols_ = tiling_->dxTilingData.cols;
43- rows_ = tiling_->dxTilingData.rows;43+ rows_ = tiling_->dxTilingData.rows;
44- binaryAddK_ = tiling_->binaryAddKDG;44+ binaryAddK_ = tiling_->binaryAddKDG;
45- colsPerCore_ = tiling_->colsPerCoreDG;45+ colsPerCore_ = tiling_->colsPerCoreDG;
46- rowsPerUB_ = tiling_->rowsPerUBDG;46+ rowsPerUB_ = tiling_->rowsPerUBDG;
47- gmOffset_ = colsPerCore_ * coreIdx_;47+ gmOffset_ = colsPerCore_ * coreIdx_;
48- colsUbLoopCount_ = CEIL_DIV(colsPerCore_, colsPerUB_);48+ colsUbLoopCount_ = CEIL_DIV(colsPerCore_, colsPerUB_);
49- dyGm_.SetGlobalBuffer((__gm__ DY_TYPE*)dy);49+ dyGm_.SetGlobalBuffer((__gm__ DY_TYPE*)dy);
50- xGm_.SetGlobalBuffer((__gm__ X_TYPE*)x);50+ xGm_.SetGlobalBuffer((__gm__ X_TYPE*)x);
51- rstdGm_.SetGlobalBuffer((__gm__ RSTD_TYPE*)rstd);51+ rstdGm_.SetGlobalBuffer((__gm__ RSTD_TYPE*)rstd);
52- dgammaGm_.SetGlobalBuffer((__gm__ float*)dgamma);52+ dgammaGm_.SetGlobalBuffer((__gm__ float*)dgamma);
53- 53+ 
54- colsPerLoopAlign_ = BLOCK_ALIGN(colsPerUB_ * sizeof(float), blockSize_) / sizeof(float);54+ colsPerLoopAlign_ = BLOCK_ALIGN(colsPerUB_ * sizeof(float), blockSize_) / sizeof(float);
55- int64_t rowsPerUbAligned = BLOCK_ALIGN(rowsPerUB_ * sizeof(float), blockSize_) / sizeof(float);55+ int64_t rowsPerUbAligned = BLOCK_ALIGN(rowsPerUB_ * sizeof(float), blockSize_) / sizeof(float);
56- Ppipe_->InitBuffer(rstdQueue_, BUFFER_NUM, (rowsPerUbAligned * sizeof(float)));56+ Ppipe_->InitBuffer(rstdQueue_, BUFFER_NUM, (rowsPerUbAligned * sizeof(float)));
57- Ppipe_->InitBuffer(dyQueue_, BUFFER_NUM, (rowsPerUB_ * colsPerLoopAlign_ * sizeof(float)));57+ Ppipe_->InitBuffer(dyQueue_, BUFFER_NUM, (rowsPerUB_ * colsPerLoopAlign_ * sizeof(float)));
58- Ppipe_->InitBuffer(xQueue_, BUFFER_NUM, (rowsPerUB_ * colsPerLoopAlign_ * sizeof(float)));58+ Ppipe_->InitBuffer(xQueue_, BUFFER_NUM, (rowsPerUB_ * colsPerLoopAlign_ * sizeof(float)));
59- Ppipe_->InitBuffer(dgammaQueue_, BUFFER_NUM, ((rowsPerUB_ + 1) * colsPerLoopAlign_ * sizeof(float)));59+ Ppipe_->InitBuffer(dgammaQueue_, BUFFER_NUM, ((rowsPerUB_ + 1) * colsPerLoopAlign_ * sizeof(float)));
60- 60+ 
61- if (!IS_FULL_LOAD) {61+ if (!IS_FULL_LOAD) {
62- Ppipe_->InitBuffer(binaryAddCacheQueue_, BUFFER_NUM,62+ Ppipe_->InitBuffer(binaryAddCacheQueue_, BUFFER_NUM,
63- ((tiling_->binaryAddKDG + RESERVESIZE) * colsPerLoopAlign_ * sizeof(float)));63+ ((tiling_->binaryAddKDG + RESERVESIZE) * colsPerLoopAlign_ * sizeof(float)));
64- Ppipe_->InitBuffer(dgammaQueue1_, BUFFER_NUM, ((rowsPerUB_ + 1) * colsPerLoopAlign_ * sizeof(float)));64+ Ppipe_->InitBuffer(dgammaQueue1_, BUFFER_NUM, ((rowsPerUB_ + 1) * colsPerLoopAlign_ * sizeof(float)));
65- }65+ }
66- }66+ }
67- 67+ 
68- __aicore__ inline void CopyInputsToUB(LocalTensor<DY_TYPE> dyLocal, LocalTensor<X_TYPE> xLocal,68+ __aicore__ inline void CopyInputsToUB(LocalTensor<DY_TYPE> dyLocal, LocalTensor<X_TYPE> xLocal,
69- LocalTensor<RSTD_TYPE> rstdLocal, int64_t inputOffset, int32_t copyLen,69+ LocalTensor<RSTD_TYPE> rstdLocal, int64_t inputOffset, int32_t copyLen,
70- int32_t curRowsNum, int32_t rstdOffset)70+ int32_t curRowsNum, int32_t rstdOffset)
71- {71+ {
72- // Datacopy Params for input_x & input_dy72+ // Datacopy Params for input_x & input_dy
73- DataCopyPadExtParams<X_TYPE> padParams_x;73+ DataCopyPadExtParams<X_TYPE> padParams_x;
74- padParams_x.isPad = false;74+ padParams_x.isPad = false;
75- DataCopyExtParams dataCopyParams_x;75+ DataCopyExtParams dataCopyParams_x;
76- dataCopyParams_x.blockCount = curRowsNum;76+ dataCopyParams_x.blockCount = curRowsNum;
77- dataCopyParams_x.blockLen = copyLen * sizeof(X_TYPE);77+ dataCopyParams_x.blockLen = copyLen * sizeof(X_TYPE);
78- dataCopyParams_x.srcStride = (cols_ - copyLen) * sizeof(X_TYPE);78+ dataCopyParams_x.srcStride = (cols_ - copyLen) * sizeof(X_TYPE);
79- dataCopyParams_x.dstStride = (colsPerUB_ - copyLen) * sizeof(X_TYPE) / blockSize_;79+ dataCopyParams_x.dstStride = (colsPerUB_ - copyLen) * sizeof(X_TYPE) / blockSize_;
80- // Datacopy Params for input_rstd80+ // Datacopy Params for input_rstd
81- DataCopyPadExtParams<RSTD_TYPE> padParams_rstd;81+ DataCopyPadExtParams<RSTD_TYPE> padParams_rstd;
82- padParams_rstd.isPad = false;82+ padParams_rstd.isPad = false;
83- DataCopyExtParams dataCopyParams_rstd;83+ DataCopyExtParams dataCopyParams_rstd;
84- dataCopyParams_rstd.blockCount = 1;84+ dataCopyParams_rstd.blockCount = 1;
85- dataCopyParams_rstd.blockLen = curRowsNum * sizeof(RSTD_TYPE);85+ dataCopyParams_rstd.blockLen = curRowsNum * sizeof(RSTD_TYPE);
86- dataCopyParams_rstd.srcStride = 0;86+ dataCopyParams_rstd.srcStride = 0;
87- dataCopyParams_rstd.dstStride = 0;87+ dataCopyParams_rstd.dstStride = 0;
88- 88+ 
89- DataCopyPad(xLocal, xGm_[inputOffset], dataCopyParams_x, padParams_x);89+ DataCopyPad(xLocal, xGm_[inputOffset], dataCopyParams_x, padParams_x);
90- DataCopyPad(dyLocal, dyGm_[inputOffset], dataCopyParams_x, padParams_x);90+ DataCopyPad(dyLocal, dyGm_[inputOffset], dataCopyParams_x, padParams_x);
91- DataCopyPad(rstdLocal, rstdGm_[rstdOffset], dataCopyParams_rstd, padParams_rstd);91+ DataCopyPad(rstdLocal, rstdGm_[rstdOffset], dataCopyParams_rstd, padParams_rstd);
92- }92+ }
93- 93+ 
94- __aicore__ inline void CopyDgammaToGm(LocalTensor<float> outLocal, uint32_t dgammaGmOffset, int32_t curCols,94+ __aicore__ inline void CopyDgammaToGm(LocalTensor<float> outLocal, uint32_t dgammaGmOffset, int32_t curCols,
95- int32_t dgammaUBOffset)95+ int32_t dgammaUBOffset)
96- {96+ {
97- DataCopyExtParams dataCopyParams;97+ DataCopyExtParams dataCopyParams;
98- dataCopyParams.blockCount = 1;98+ dataCopyParams.blockCount = 1;
99- dataCopyParams.blockLen = curCols * sizeof(float);99+ dataCopyParams.blockLen = curCols * sizeof(float);
100- dataCopyParams.srcStride = 0;100+ dataCopyParams.srcStride = 0;
101- dataCopyParams.dstStride = 0;101+ dataCopyParams.dstStride = 0;
102- 102+ 
103- DataCopyPad(dgammaGm_[dgammaGmOffset], outLocal[dgammaUBOffset], dataCopyParams);103+ DataCopyPad(dgammaGm_[dgammaGmOffset], outLocal[dgammaUBOffset], dataCopyParams);
104- }104+ }
105- 105+ 
106- __aicore__ inline void VFCalcPreDgamma(__local_mem__ DY_TYPE* dyAddr, __local_mem__ X_TYPE* xAddr,106+ __aicore__ inline void VFCalcPreDgamma(__ubuf__ DY_TYPE* dyAddr, __ubuf__ X_TYPE* xAddr,
107- __local_mem__ RSTD_TYPE* rstdAddr, __local_mem__ float* dgammaOutAddr,107+ __ubuf__ RSTD_TYPE* rstdAddr, __ubuf__ float* dgammaOutAddr,
108- uint16_t curUBLoopColsCount, int32_t curRowsNum)108+ uint16_t curUBLoopColsCount, int32_t curRowsNum)
109- {109+ {
110- uint16_t colsRegLoopCount = CEIL_DIV(curUBLoopColsCount, vlFp32_);110+ uint16_t colsRegLoopCount = CEIL_DIV(curUBLoopColsCount, vlFp32_);
111- uint32_t colsPerLoop = colsPerUB_;111+ uint32_t colsPerLoop = colsPerUB_;
112- uint32_t colsPerLoopAlign = colsPerLoopAlign_;112+ uint32_t colsPerLoopAlign = colsPerLoopAlign_;
113- 113+ 
114- __VEC_SCOPE__114+ __VEC_SCOPE__
115- {115+ {
116- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();116+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
117- // 计算二分累加前的乘法计算117+ // 计算二分累加前的乘法计算
118- for (uint16_t k = 0; k < static_cast<uint16_t>(curRowsNum); k++) {118+ for (uint16_t k = 0; k < static_cast<uint16_t>(curRowsNum); k++) {
119- CalcMulRes<DY_TYPE, X_TYPE, RSTD_TYPE>(dyAddr, xAddr, rstdAddr, dgammaOutAddr, pregMain,119+ CalcMulRes<DY_TYPE, X_TYPE, RSTD_TYPE>(dyAddr, xAddr, rstdAddr, dgammaOutAddr, pregMain,
120- (k * colsRegLoopCount) * vlFp32_, k);120+ (k * colsRegLoopCount) * vlFp32_, k);
121- }121+ }
122- }122+ }
123- }123+ }
124- 124+ 
125- __aicore__ inline void VFDuplicateRows(__local_mem__ float* srcAddr, uint32_t padRowsLen, uint64_t rowsBoundLine)125+ __aicore__ inline void VFDuplicateRows(__ubuf__ float* srcAddr, uint32_t padRowsLen, uint64_t rowsBoundLine)
126- {126+ {
127- __VEC_SCOPE__127+ __VEC_SCOPE__
128- {128+ {
129- RegTensor<float> tempReg;129+ RegTensor<float> tempReg;
130- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();130+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
131- uint32_t sreg0 = padRowsLen;131+ uint32_t sreg0 = padRowsLen;
132- MaskReg pregLoop = UpdateMask<float>(sreg0);132+ MaskReg pregLoop = UpdateMask<float>(sreg0);
133- // 填充数据133+ // 填充数据
134- AscendC::MicroAPI::Duplicate(tempReg, 0);134+ AscendC::MicroAPI::Duplicate(tempReg, 0);
135- AscendC::MicroAPI::DataCopy(srcAddr + static_cast<uint32_t>(rowsBoundLine), tempReg, pregLoop);135+ AscendC::MicroAPI::StoreAlign(srcAddr + static_cast<uint32_t>(rowsBoundLine), tempReg, pregLoop);
136- }136+ }
137- }137+ }
138- 138+ 
139- __aicore__ inline void VFBinaryReduceSumWithoutTail(__local_mem__ float* dgammaOutAddr, uint16_t curUbLoopColsCount,139+ __aicore__ inline void VFBinaryReduceSumWithoutTail(__ubuf__ float* dgammaOutAddr, uint16_t curUbLoopColsCount,
140- int64_t rows)140+ int64_t rows)
141- {141+ {
142- uint32_t BinaryAddNumLevel2 = rows / REDUCEBY8ELENUM;142+ uint32_t BinaryAddNumLevel2 = rows / REDUCEBY8ELENUM;
143- uint32_t BinaryAddNumLevel1 = BinaryAddNumLevel2 <= REDUCEBY8ELENUM ? 1 : BinaryAddNumLevel2 / 16;143+ uint32_t BinaryAddNumLevel1 = BinaryAddNumLevel2 <= REDUCEBY8ELENUM ? 1 : BinaryAddNumLevel2 / 16;
144- if (rows <= REDUCEBY8ELENUM) {144+ if (rows <= REDUCEBY8ELENUM) {
145- BinaryAddNumLevel1 = 0;145+ BinaryAddNumLevel1 = 0;
146- }146+ }
147- __VEC_SCOPE__147+ __VEC_SCOPE__
148- {148+ {
149- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();149+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
150- // 计算二分累加level2150+ // 计算二分累加level2
151- for (uint16_t i = 0; i < static_cast<uint16_t>(BinaryAddNumLevel2); i++) {151+ for (uint16_t i = 0; i < static_cast<uint16_t>(BinaryAddNumLevel2); i++) {
152- reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, i * vlFp32_);152+ reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, i * vlFp32_);
153- }153+ }
154- MicroAPI::LocalMemBar<MicroAPI::MemType::VEC_STORE, MicroAPI::MemType::VEC_LOAD>();154+ MicroAPI::LocalMemBar<MicroAPI::MemType::VEC_STORE, MicroAPI::MemType::VEC_LOAD>();
155- // 计算二分累加level1155+ // 计算二分累加level1
156- if (BinaryAddNumLevel1 == 1) {156+ if (BinaryAddNumLevel1 == 1) {
157- if (BinaryAddNumLevel2 == REDUCEBY8ELENUM) {157+ if (BinaryAddNumLevel2 == REDUCEBY8ELENUM) {
158- reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, 0);158+ reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, 0);
159- } else if (BinaryAddNumLevel2 == REDUCEBY4ELENUM) {159+ } else if (BinaryAddNumLevel2 == REDUCEBY4ELENUM) {
160- reduceSumCompressedBy4(dgammaOutAddr, pregMain, vlFp32_, 0);160+ reduceSumCompressedBy4(dgammaOutAddr, pregMain, vlFp32_, 0);
161- } else if (BinaryAddNumLevel2 == REDUCEBY2ELENUM) {161+ } else if (BinaryAddNumLevel2 == REDUCEBY2ELENUM) {
162- reduceSumCompressedBy2(dgammaOutAddr, pregMain, vlFp32_, 0);162+ reduceSumCompressedBy2(dgammaOutAddr, pregMain, vlFp32_, 0);
163- } else if (BinaryAddNumLevel2 == REDUCEBY1ELENUM) {163+ } else if (BinaryAddNumLevel2 == REDUCEBY1ELENUM) {
164- reduceSumCompressedBy1(dgammaOutAddr, pregMain, vlFp32_);164+ reduceSumCompressedBy1(dgammaOutAddr, pregMain, vlFp32_);
165- }165+ }
166- } else {166+ } else {
167- for (uint16_t i = 0; i < BinaryAddNumLevel1; i++) {167+ for (uint16_t i = 0; i < BinaryAddNumLevel1; i++) {
168- reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, i * vlFp32_);168+ reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, i * vlFp32_);
169- }169+ }
170- // 计算二分累加level0170+ // 计算二分累加level0
171- if (BinaryAddNumLevel1 == REDUCEBY8ELENUM) {171+ if (BinaryAddNumLevel1 == REDUCEBY8ELENUM) {
172- reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, 0);172+ reduceSumCompressedBy8(dgammaOutAddr, pregMain, vlFp32_, 0);
173- } else if (BinaryAddNumLevel1 == REDUCEBY4ELENUM) {173+ } else if (BinaryAddNumLevel1 == REDUCEBY4ELENUM) {
174- reduceSumCompressedBy4(dgammaOutAddr, pregMain, vlFp32_, 0);174+ reduceSumCompressedBy4(dgammaOutAddr, pregMain, vlFp32_, 0);
175- } else if (BinaryAddNumLevel1 == REDUCEBY2ELENUM) {175+ } else if (BinaryAddNumLevel1 == REDUCEBY2ELENUM) {
176- reduceSumCompressedBy2(dgammaOutAddr, pregMain, vlFp32_, 0);176+ reduceSumCompressedBy2(dgammaOutAddr, pregMain, vlFp32_, 0);
177- } else if (BinaryAddNumLevel1 == REDUCEBY1ELENUM) {177+ } else if (BinaryAddNumLevel1 == REDUCEBY1ELENUM) {
178- reduceSumCompressedBy1(dgammaOutAddr, pregMain, vlFp32_);178+ reduceSumCompressedBy1(dgammaOutAddr, pregMain, vlFp32_);
179- }179+ }
180- }180+ }
181- if (rows == REDUCEBY4ELENUM) {181+ if (rows == REDUCEBY4ELENUM) {
182- reduceSumCompressedBy4(dgammaOutAddr, pregMain, vlFp32_, 0);182+ reduceSumCompressedBy4(dgammaOutAddr, pregMain, vlFp32_, 0);
183- } else if (rows == REDUCEBY2ELENUM) {183+ } else if (rows == REDUCEBY2ELENUM) {
184- reduceSumCompressedBy2(dgammaOutAddr, pregMain, vlFp32_, 0);184+ reduceSumCompressedBy2(dgammaOutAddr, pregMain, vlFp32_, 0);
185- } else if (rows == REDUCEBY1ELENUM) {185+ } else if (rows == REDUCEBY1ELENUM) {
186- reduceSumCompressedBy1(dgammaOutAddr, pregMain, vlFp32_);186+ reduceSumCompressedBy1(dgammaOutAddr, pregMain, vlFp32_);
187- }187+ }
188- }188+ }
189- }189+ }
190- 190+ 
191- __aicore__ inline void VFHandleTailRows(__local_mem__ float* dgammaOutAddr, uint16_t rowsTail,191+ __aicore__ inline void VFHandleTailRows(__ubuf__ float* dgammaOutAddr, uint16_t rowsTail, uint64_t tailDataOffset)
192- uint64_t tailDataOffset)192+ {
193- {193+ uint32_t BinaryAddTailNum = (rowsTail + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;
194- uint32_t BinaryAddTailNum = (rowsTail + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;194+ 
195- 195+ __VEC_SCOPE__
196- __VEC_SCOPE__196+ {
197- {197+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
198- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();198+ uint64_t rowsBoundLine = rows_ * vlFp32_;
199- uint64_t rowsBoundLine = rows_ * vlFp32_;199+ for (uint16_t i = 0; i < static_cast<uint16_t>(BinaryAddTailNum - 1); i++) {
200- for (uint16_t i = 0; i < static_cast<uint16_t>(BinaryAddTailNum - 1); i++) {200+ reduceSumCompressedBy8WithOutPad(dgammaOutAddr, dgammaOutAddr + tailDataOffset, pregMain,
201- reduceSumCompressedBy8WithOutPad(dgammaOutAddr, dgammaOutAddr + tailDataOffset, pregMain,201+ i * vlFp32_ * COMPRESSBY8ELENUM, vlFp32_);
202- i * vlFp32_ * COMPRESSBY8ELENUM, vlFp32_);202+ }
203- }203+ reduceSumCompressedBy8WithPad(dgammaOutAddr, dgammaOutAddr, pregMain,
204- reduceSumCompressedBy8WithPad(dgammaOutAddr, dgammaOutAddr, pregMain,204+ (BinaryAddTailNum - 1) * vlFp32_ * COMPRESSBY8ELENUM, rowsBoundLine, vlFp32_,
205- (BinaryAddTailNum - 1) * vlFp32_ * COMPRESSBY8ELENUM, rowsBoundLine, vlFp32_,205+ tailDataOffset);
206- tailDataOffset);206+ }
207- }207+ }
208- }208+ 
209- 209+ __aicore__ inline void VFHandleTailRowsWithTwoBuffer(__ubuf__ float* dgammaOutAddr, __ubuf__ float* dgammaOutAddr1,
210- __aicore__ inline void VFHandleTailRowsWithTwoBuffer(__local_mem__ float* dgammaOutAddr,210+ uint64_t tailRowsNum)
211- __local_mem__ float* dgammaOutAddr1, uint64_t tailRowsNum)211+ {
212- {212+ uint32_t BinaryAddTailNum = (tailRowsNum + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;
213- uint32_t BinaryAddTailNum = (tailRowsNum + COMPRESSBY8ELENUM - 1) / COMPRESSBY8ELENUM;213+ uint32_t tailDataOffset = 0;
214- uint32_t tailDataOffset = 0;214+ 
215- 215+ __VEC_SCOPE__
216- __VEC_SCOPE__216+ {
217- {217+ MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();
218- MaskReg pregMain = CreateMask<float, MaskPattern::ALL>();218+ uint64_t rowsBoundLine = tailRowsNum * vlFp32_;
219- uint64_t rowsBoundLine = tailRowsNum * vlFp32_;219+ for (uint16_t i = 0; i < static_cast<uint16_t>(BinaryAddTailNum - 1); i++) {
220- for (uint16_t i = 0; i < static_cast<uint16_t>(BinaryAddTailNum - 1); i++) {220+ reduceSumCompressedBy8WithOutPad(dgammaOutAddr, dgammaOutAddr1, pregMain,
221- reduceSumCompressedBy8WithOutPad(dgammaOutAddr, dgammaOutAddr1, pregMain,221+ i * vlFp32_ * COMPRESSBY8ELENUM, vlFp32_);
222- i * vlFp32_ * COMPRESSBY8ELENUM, vlFp32_);222+ }
223- }223+ reduceSumCompressedBy8WithPad(dgammaOutAddr, dgammaOutAddr1, pregMain,
224- reduceSumCompressedBy8WithPad(dgammaOutAddr, dgammaOutAddr1, pregMain,224+ (BinaryAddTailNum - 1) * vlFp32_ * COMPRESSBY8ELENUM, rowsBoundLine, vlFp32_,
225- (BinaryAddTailNum - 1) * vlFp32_ * COMPRESSBY8ELENUM, rowsBoundLine, vlFp32_,225+ tailDataOffset);
226- tailDataOffset);226+ }
227- }227+ }
228- }228+ 
229- 229+ __aicore__ inline void CalcDgamma(uint32_t inputOffset, uint32_t currentCols, bool isWithPad)
230- __aicore__ inline void CalcDgamma(uint32_t inputOffset, uint32_t currentCols, bool isWithPad)230+ {
231- {231+ LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();
232- LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();232+ LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();
233- LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();233+ LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();
234- LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();234+ LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();
235- LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();235+ 
236- 236+ CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, 0);
237- CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, 0);237+ xQueue_.EnQue(xLocal);
238- xQueue_.EnQue(xLocal);238+ rstdQueue_.EnQue(rstdLocal);
239- rstdQueue_.EnQue(rstdLocal);239+ dyQueue_.EnQue(dyLocal);
240- dyQueue_.EnQue(dyLocal);240+ 
241- 241+ dyLocal = dyQueue_.template DeQue<DY_TYPE>();
242- dyLocal = dyQueue_.template DeQue<DY_TYPE>();242+ xLocal = xQueue_.template DeQue<X_TYPE>();
243- xLocal = xQueue_.template DeQue<X_TYPE>();243+ rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
244- rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();244+ 
245- 245+ __ubuf__ DY_TYPE* dyAddr = (__ubuf__ DY_TYPE*)dyLocal[0].GetPhyAddr();
246- __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();246+ __ubuf__ X_TYPE* xAddr = (__ubuf__ X_TYPE*)xLocal[0].GetPhyAddr();
247- __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();247+ __ubuf__ RSTD_TYPE* rstdAddr = (__ubuf__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
248- __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();248+ __ubuf__ float* dgammaOutAddr = (__ubuf__ float*)dgammaOutLocal[0].GetPhyAddr();
249- __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();249+ 
250- 250+ VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
251- VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);251+ dyQueue_.FreeTensor(dyLocal);
252- dyQueue_.FreeTensor(dyLocal);252+ xQueue_.FreeTensor(xLocal);
253- xQueue_.FreeTensor(xLocal);253+ rstdQueue_.FreeTensor(rstdLocal);
254- rstdQueue_.FreeTensor(rstdLocal);254+ if (isWithPad) {
255- if (isWithPad) {255+ uint32_t mainRows = rows_ - tiling_->rowsTailDG;
256- uint32_t mainRows = rows_ - tiling_->rowsTailDG;256+ VFDuplicateRows(dgammaOutAddr, vlFp32_, rows_ * vlFp32_);
257- VFDuplicateRows(dgammaOutAddr, vlFp32_, rows_ * vlFp32_);257+ tailDataOffset_ = mainRows * vlFp32_;
258- tailDataOffset_ = mainRows * vlFp32_;258+ VFHandleTailRows(dgammaOutAddr, tiling_->rowsTailDG, tailDataOffset_);
259- VFHandleTailRows(dgammaOutAddr, tiling_->rowsTailDG, tailDataOffset_);259+ VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, mainRows);
260- VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, mainRows);260+ } else {
261- } else {261+ VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rows_);
262- VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rows_);262+ }
263- }263+ 
264- 264+ dgammaQueue_.EnQue(dgammaOutLocal);
265- dgammaQueue_.EnQue(dgammaOutLocal);265+ dgammaOutLocal = dgammaQueue_.template DeQue<float>();
266- dgammaOutLocal = dgammaQueue_.template DeQue<float>();266+ CopyDgammaToGm(dgammaOutLocal, inputOffset, currentCols, 0);
267- CopyDgammaToGm(dgammaOutLocal, inputOffset, currentCols, 0);267+ 
268- 268+ dgammaQueue_.FreeTensor(dgammaOutLocal);
269- dgammaQueue_.FreeTensor(dgammaOutLocal);269+ }
270- }270+ 
271- 271+ __aicore__ inline void CalcLargeRowsDgamma(uint32_t inputOffset, uint32_t currentCols, uint32_t i,
272- __aicore__ inline void CalcLargeRowsDgamma(uint32_t inputOffset, uint32_t currentCols, uint32_t i,272+ LocalTensor<float> binaryAddCacheLocal)
273- LocalTensor<float> binaryAddCacheLocal)273+ {
274- {274+ LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();
275- LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();275+ LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();
276- LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();276+ LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();
277- LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();277+ LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();
278- LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();278+ int64_t cacheID = GetCacheID(i);
279- int64_t cacheID = GetCacheID(i);279+ uint32_t rstdOffset = i * rowsPerUB_;
280- uint32_t rstdOffset = i * rowsPerUB_;280+ 
281- 281+ CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);
282- CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);282+ xQueue_.EnQue(xLocal);
283- xQueue_.EnQue(xLocal);283+ rstdQueue_.EnQue(rstdLocal);
284- rstdQueue_.EnQue(rstdLocal);284+ dyQueue_.EnQue(dyLocal);
285- dyQueue_.EnQue(dyLocal);285+ 
286- 286+ dyLocal = dyQueue_.template DeQue<DY_TYPE>();
287- dyLocal = dyQueue_.template DeQue<DY_TYPE>();287+ xLocal = xQueue_.template DeQue<X_TYPE>();
288- xLocal = xQueue_.template DeQue<X_TYPE>();288+ rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
289- rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();289+ 
290- 290+ __ubuf__ DY_TYPE* dyAddr = (__ubuf__ DY_TYPE*)dyLocal[0].GetPhyAddr();
291- __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();291+ __ubuf__ X_TYPE* xAddr = (__ubuf__ X_TYPE*)xLocal[0].GetPhyAddr();
292- __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();292+ __ubuf__ RSTD_TYPE* rstdAddr = (__ubuf__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
293- __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();293+ __ubuf__ float* dgammaOutAddr = (__ubuf__ float*)dgammaOutLocal[0].GetPhyAddr();
294- __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();294+ 
295- 295+ VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
296- VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);296+ dyQueue_.FreeTensor(dyLocal);
297- dyQueue_.FreeTensor(dyLocal);297+ xQueue_.FreeTensor(xLocal);
298- xQueue_.FreeTensor(xLocal);298+ rstdQueue_.FreeTensor(rstdLocal);
299- rstdQueue_.FreeTensor(rstdLocal);299+ 
300- 300+ VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rowsPerUB_);
301- VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rowsPerUB_);301+ UpdateCache(binaryAddCacheLocal, dgammaOutAddr, cacheID, vlFp32_);
302- UpdateCache(binaryAddCacheLocal, dgammaOutAddr, cacheID, vlFp32_);302+ dgammaQueue_.FreeTensor(dgammaOutLocal);
303- dgammaQueue_.FreeTensor(dgammaOutLocal);303+ }
304- }304+ 
305- 305+ __aicore__ inline void CalcLargeRowsDgammaWithPad(uint32_t inputOffset, uint32_t tailDyXOffset,
306- __aicore__ inline void CalcLargeRowsDgammaWithPad(uint32_t inputOffset, uint32_t tailDyXOffset,306+ uint32_t tailRstdOffset, uint32_t currentCols, uint32_t i,
307- uint32_t tailRstdOffset, uint32_t currentCols, uint32_t i,307+ LocalTensor<float> binaryAddCacheLocal)
308- LocalTensor<float> binaryAddCacheLocal)308+ {
309- {309+ LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();
310- LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();310+ LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();
311- LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();311+ LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();
312- LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();312+ LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();
313- LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();313+ int64_t cacheID = GetCacheID(i);
314- int64_t cacheID = GetCacheID(i);314+ uint32_t rstdOffset = i * rowsPerUB_;
315- uint32_t rstdOffset = i * rowsPerUB_;315+ 
316- 316+ CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);
317- CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);317+ xQueue_.EnQue(xLocal);
318- xQueue_.EnQue(xLocal);318+ rstdQueue_.EnQue(rstdLocal);
319- rstdQueue_.EnQue(rstdLocal);319+ dyQueue_.EnQue(dyLocal);
320- dyQueue_.EnQue(dyLocal);320+ 
321- 321+ dyLocal = dyQueue_.template DeQue<DY_TYPE>();
322- dyLocal = dyQueue_.template DeQue<DY_TYPE>();322+ xLocal = xQueue_.template DeQue<X_TYPE>();
323- xLocal = xQueue_.template DeQue<X_TYPE>();323+ rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
324- rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();324+ 
325- 325+ __ubuf__ DY_TYPE* dyAddr = (__ubuf__ DY_TYPE*)dyLocal[0].GetPhyAddr();
326- __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();326+ __ubuf__ X_TYPE* xAddr = (__ubuf__ X_TYPE*)xLocal[0].GetPhyAddr();
327- __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();327+ __ubuf__ RSTD_TYPE* rstdAddr = (__ubuf__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
328- __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();328+ __ubuf__ float* dgammaOutAddr = (__ubuf__ float*)dgammaOutLocal[0].GetPhyAddr();
329- __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();329+ 
330- 330+ VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
331- VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);331+ dyQueue_.FreeTensor(dyLocal);
332- dyQueue_.FreeTensor(dyLocal);332+ xQueue_.FreeTensor(xLocal);
333- xQueue_.FreeTensor(xLocal);333+ rstdQueue_.FreeTensor(rstdLocal);
334- rstdQueue_.FreeTensor(rstdLocal);334+ 
335- 335+ // 处理累加的尾块
336- // 处理累加的尾块336+ LocalTensor<RSTD_TYPE> rstdLocal1 = rstdQueue_.template AllocTensor<RSTD_TYPE>();
337- LocalTensor<RSTD_TYPE> rstdLocal1 = rstdQueue_.template AllocTensor<RSTD_TYPE>();337+ LocalTensor<DY_TYPE> dyLocal1 = dyQueue_.template AllocTensor<DY_TYPE>();
338- LocalTensor<DY_TYPE> dyLocal1 = dyQueue_.template AllocTensor<DY_TYPE>();338+ LocalTensor<X_TYPE> xLocal1 = xQueue_.template AllocTensor<X_TYPE>();
339- LocalTensor<X_TYPE> xLocal1 = xQueue_.template AllocTensor<X_TYPE>();339+ LocalTensor<float> dgammaOutLocal1 = dgammaQueue1_.template AllocTensor<float>();
340- LocalTensor<float> dgammaOutLocal1 = dgammaQueue1_.template AllocTensor<float>();340+ int32_t tailRowsNum = i == tiling_->tailBlockCountWithoutPadDG ? rows_ % rowsPerUB_ : rowsPerUB_;
341- int32_t tailRowsNum = i == tiling_->tailBlockCountWithoutPadDG ? rows_ % rowsPerUB_ : rowsPerUB_;341+ CopyInputsToUB(dyLocal1, xLocal1, rstdLocal1, tailDyXOffset, currentCols, tailRowsNum, tailRstdOffset);
342- CopyInputsToUB(dyLocal1, xLocal1, rstdLocal1, tailDyXOffset, currentCols, tailRowsNum, tailRstdOffset);342+ xQueue_.EnQue(xLocal1);
343- xQueue_.EnQue(xLocal1);343+ rstdQueue_.EnQue(rstdLocal1);
344- rstdQueue_.EnQue(rstdLocal1);344+ dyQueue_.EnQue(dyLocal1);
345- dyQueue_.EnQue(dyLocal1);345+ 
346- 346+ dyLocal1 = dyQueue_.template DeQue<DY_TYPE>();
347- dyLocal1 = dyQueue_.template DeQue<DY_TYPE>();347+ xLocal1 = xQueue_.template DeQue<X_TYPE>();
348- xLocal1 = xQueue_.template DeQue<X_TYPE>();348+ rstdLocal1 = rstdQueue_.template DeQue<RSTD_TYPE>();
349- rstdLocal1 = rstdQueue_.template DeQue<RSTD_TYPE>();349+ 
350- 350+ __ubuf__ DY_TYPE* dyAddr1 = (__ubuf__ DY_TYPE*)dyLocal1[0].GetPhyAddr();
351- __local_mem__ DY_TYPE* dyAddr1 = (__local_mem__ DY_TYPE*)dyLocal1[0].GetPhyAddr();351+ __ubuf__ X_TYPE* xAddr1 = (__ubuf__ X_TYPE*)xLocal1[0].GetPhyAddr();
352- __local_mem__ X_TYPE* xAddr1 = (__local_mem__ X_TYPE*)xLocal1[0].GetPhyAddr();352+ __ubuf__ RSTD_TYPE* rstdAddr1 = (__ubuf__ RSTD_TYPE*)rstdLocal1[0].GetPhyAddr();
353- __local_mem__ RSTD_TYPE* rstdAddr1 = (__local_mem__ RSTD_TYPE*)rstdLocal1[0].GetPhyAddr();353+ __ubuf__ float* dgammaOutAddr1 = (__ubuf__ float*)dgammaOutLocal1[0].GetPhyAddr();
354- __local_mem__ float* dgammaOutAddr1 = (__local_mem__ float*)dgammaOutLocal1[0].GetPhyAddr();354+ 
355- 355+ VFCalcPreDgamma(dyAddr1, xAddr1, rstdAddr1, dgammaOutAddr1, currentCols, rowsPerUB_);
356- VFCalcPreDgamma(dyAddr1, xAddr1, rstdAddr1, dgammaOutAddr1, currentCols, rowsPerUB_);356+ dyQueue_.FreeTensor(dyLocal1);
357- dyQueue_.FreeTensor(dyLocal1);357+ xQueue_.FreeTensor(xLocal1);
358- xQueue_.FreeTensor(xLocal1);358+ rstdQueue_.FreeTensor(rstdLocal1);
359- rstdQueue_.FreeTensor(rstdLocal1);359+ 
360- 360+ if (i == tiling_->tailBlockCountWithoutPadDG) {
361- if (i == tiling_->tailBlockCountWithoutPadDG) {361+ VFDuplicateRows(dgammaOutAddr1, vlFp32_, tailRowsNum * vlFp32_);
362- VFDuplicateRows(dgammaOutAddr1, vlFp32_, tailRowsNum * vlFp32_);362+ VFHandleTailRowsWithTwoBuffer(dgammaOutAddr, dgammaOutAddr1, tailRowsNum);
363- VFHandleTailRowsWithTwoBuffer(dgammaOutAddr, dgammaOutAddr1, tailRowsNum);363+ } else {
364- } else {364+ VFHandleTailRowsWithTwoBuffer(dgammaOutAddr, dgammaOutAddr1, rowsPerUB_);
365- VFHandleTailRowsWithTwoBuffer(dgammaOutAddr, dgammaOutAddr1, rowsPerUB_);365+ }
366- }366+ dgammaQueue1_.FreeTensor(dgammaOutLocal1);
367- dgammaQueue1_.FreeTensor(dgammaOutLocal1);367+ 
368- 368+ VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rowsPerUB_);
369- VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rowsPerUB_);369+ UpdateCache(binaryAddCacheLocal, dgammaOutAddr, cacheID, vlFp32_);
370- UpdateCache(binaryAddCacheLocal, dgammaOutAddr, cacheID, vlFp32_);370+ dgammaQueue_.FreeTensor(dgammaOutLocal);
371- dgammaQueue_.FreeTensor(dgammaOutLocal);371+ }
372- }372+ 
373- 373+ __aicore__ inline void HandlingLargeRows2KAlign(uint32_t startOffset, uint32_t currentCols)
374- __aicore__ inline void HandlingLargeRows2KAlign(uint32_t startOffset, uint32_t currentCols)374+ {
375- {375+ LocalTensor<float> binaryAddCacheLocal = binaryAddCacheQueue_.template AllocTensor<float>();
376- LocalTensor<float> binaryAddCacheLocal = binaryAddCacheQueue_.template AllocTensor<float>();376+ for (int32_t i = 0; i < tiling_->mainBlockCountDG; i++) {
377- for (int32_t i = 0; i < tiling_->mainBlockCountDG; i++) {377+ uint32_t inputOffset = startOffset + i * cols_ * rowsPerUB_;
378- uint32_t inputOffset = startOffset + i * cols_ * rowsPerUB_;378+ CalcLargeRowsDgamma(inputOffset, currentCols, i, binaryAddCacheLocal);
379- CalcLargeRowsDgamma(inputOffset, currentCols, i, binaryAddCacheLocal);379+ }
380- }380+ binaryAddCacheQueue_.EnQue(binaryAddCacheLocal);
381- binaryAddCacheQueue_.EnQue(binaryAddCacheLocal);381+ binaryAddCacheLocal = binaryAddCacheQueue_.template DeQue<float>();
382- binaryAddCacheLocal = binaryAddCacheQueue_.template DeQue<float>();382+ CopyDgammaToGm(binaryAddCacheLocal, startOffset, currentCols, binaryAddK_ * vlFp32_);
383- CopyDgammaToGm(binaryAddCacheLocal, startOffset, currentCols, binaryAddK_ * vlFp32_);383+ binaryAddCacheQueue_.FreeTensor(binaryAddCacheLocal);
384- binaryAddCacheQueue_.FreeTensor(binaryAddCacheLocal);384+ }
385- }385+ 
386- 386+ __aicore__ inline void HandlingLargeRows2KUnAlign(uint32_t startOffset, uint32_t currentCols)
387- __aicore__ inline void HandlingLargeRows2KUnAlign(uint32_t startOffset, uint32_t currentCols)387+ {
388- {388+ LocalTensor<float> binaryAddCacheLocal = binaryAddCacheQueue_.template AllocTensor<float>();
389- LocalTensor<float> binaryAddCacheLocal = binaryAddCacheQueue_.template AllocTensor<float>();389+ // 处理row对齐尾块
390- // 处理row对齐尾块390+ for (uint32_t i = 0; i < tiling_->tailBlockCountWithoutPadDG; i++) {
391- for (uint32_t i = 0; i < tiling_->tailBlockCountWithoutPadDG; i++) {391+ uint32_t loopOffset = i * cols_ * rowsPerUB_;
392- uint32_t loopOffset = i * cols_ * rowsPerUB_;392+ uint32_t inputOffset = startOffset + loopOffset;
393- uint32_t inputOffset = startOffset + loopOffset;393+ uint32_t tailDyXOffset = inputOffset + tiling_->powerOfTwoBlockCountDG * rowsPerUB_ * cols_;
394- uint32_t tailDyXOffset = inputOffset + tiling_->powerOfTwoBlockCountDG * rowsPerUB_ * cols_;394+ uint32_t tailRstdOffset = i * rowsPerUB_ + tiling_->powerOfTwoBlockCountDG * rowsPerUB_;
395- uint32_t tailRstdOffset = i * rowsPerUB_ + tiling_->powerOfTwoBlockCountDG * rowsPerUB_;395+ CalcLargeRowsDgammaWithPad(inputOffset, tailDyXOffset, tailRstdOffset, currentCols, i, binaryAddCacheLocal);
396- CalcLargeRowsDgammaWithPad(inputOffset, tailDyXOffset, tailRstdOffset, currentCols, i, binaryAddCacheLocal);396+ }
397- }397+ // 处理row不对齐尾块
398- // 处理row不对齐尾块398+ for (uint32_t i = tiling_->tailBlockCountWithoutPadDG;
399- for (uint32_t i = tiling_->tailBlockCountWithoutPadDG;399+ i < tiling_->tailBlockCountWithoutPadDG + tiling_->tailBlockCountwithPadDG; i++) {
400- i < tiling_->tailBlockCountWithoutPadDG + tiling_->tailBlockCountwithPadDG; i++) {400+ uint32_t loopOffset = i * cols_ * rowsPerUB_;
401- uint32_t loopOffset = i * cols_ * rowsPerUB_;401+ uint32_t inputOffset = startOffset + loopOffset;
402- uint32_t inputOffset = startOffset + loopOffset;402+ uint32_t tailDyXOffset = startOffset + tiling_->mainBlockCountDG * rowsPerUB_ * cols_;
403- uint32_t tailDyXOffset = startOffset + tiling_->mainBlockCountDG * rowsPerUB_ * cols_;403+ uint32_t tailRstdOffset = tiling_->mainBlockCountDG * rowsPerUB_;
404- uint32_t tailRstdOffset = tiling_->mainBlockCountDG * rowsPerUB_;404+ CalcLargeRowsDgammaWithPad(inputOffset, tailDyXOffset, tailRstdOffset, currentCols, i, binaryAddCacheLocal);
405- CalcLargeRowsDgammaWithPad(inputOffset, tailDyXOffset, tailRstdOffset, currentCols, i, binaryAddCacheLocal);405+ }
406- }406+ // 处理正常主块
407- // 处理正常主块407+ for (uint32_t i = tiling_->tailBlockCountWithoutPadDG + tiling_->tailBlockCountwithPadDG;
408- for (uint32_t i = tiling_->tailBlockCountWithoutPadDG + tiling_->tailBlockCountwithPadDG;408+ i < tiling_->powerOfTwoBlockCountDG; i++) {
409- i < tiling_->powerOfTwoBlockCountDG; i++) {409+ uint32_t loopOffset = i * cols_ * rowsPerUB_;
410- uint32_t loopOffset = i * cols_ * rowsPerUB_;410+ uint32_t inputOffset = startOffset + loopOffset;
411- uint32_t inputOffset = startOffset + loopOffset;411+ CalcLargeRowsDgamma(inputOffset, currentCols, i, binaryAddCacheLocal);
412- CalcLargeRowsDgamma(inputOffset, currentCols, i, binaryAddCacheLocal);412+ }
413- }413+ binaryAddCacheQueue_.EnQue(binaryAddCacheLocal);
414- binaryAddCacheQueue_.EnQue(binaryAddCacheLocal);414+ binaryAddCacheLocal = binaryAddCacheQueue_.template DeQue<float>();
415- binaryAddCacheLocal = binaryAddCacheQueue_.template DeQue<float>();415+ CopyDgammaToGm(binaryAddCacheLocal, startOffset, currentCols, binaryAddK_ * vlFp32_);
416- CopyDgammaToGm(binaryAddCacheLocal, startOffset, currentCols, binaryAddK_ * vlFp32_);416+ binaryAddCacheQueue_.FreeTensor(binaryAddCacheLocal);
417- binaryAddCacheQueue_.FreeTensor(binaryAddCacheLocal);417+ }
418- }418+ 
419- 419+ __aicore__ inline void Process()
420- __aicore__ inline void Process()420+ {
421- {421+ if (coreIdx_ >= tiling_->usedCoreNumDG) {
422- if (coreIdx_ >= tiling_->usedCoreNumDG) {422+ return;
423- return;423+ }
424- }424+ int64_t inputOffset = 0;
425- int64_t inputOffset = 0;425+ curCols_ = colsPerUB_;
426- curCols_ = colsPerUB_;426+ if ((coreIdx_ == tiling_->usedCoreNumDG - 1) && !tiling_->isMultiColset) {
427- if ((coreIdx_ == tiling_->usedCoreNumDG - 1) && !tiling_->isMultiColset) {427+ curCols_ = tiling_->colsPerTailCoreDG;
428- curCols_ = tiling_->colsPerTailCoreDG;428+ }
429- }429+ // main core calc
430- // main core calc430+ bool isWithPad = tiling_->rowsTailDG == 0 ? false : true;
431- bool isWithPad = tiling_->rowsTailDG == 0 ? false : true;431+ for (int64_t curLoop = 0; curLoop < colsUbLoopCount_; curLoop++) {
432- for (int64_t curLoop = 0; curLoop < colsUbLoopCount_; curLoop++) {432+ inputOffset = curLoop * colsPerUB_ + gmOffset_;
433- inputOffset = curLoop * colsPerUB_ + gmOffset_;433+ CalcDgamma(inputOffset, curCols_, isWithPad);
434- CalcDgamma(inputOffset, curCols_, isWithPad);434+ }
435- }435+ // Handle tail
436- // Handle tail436+ inputOffset = tiling_->usedCoreNumDG * colsPerCore_ + coreIdx_ * vlFp32_;
437- inputOffset = tiling_->usedCoreNumDG * colsPerCore_ + coreIdx_ * vlFp32_;437+ if (coreIdx_ < tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {
438- if (coreIdx_ < tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {438+ CalcDgamma(inputOffset, curCols_, isWithPad);
439- CalcDgamma(inputOffset, curCols_, isWithPad);439+ } else if (coreIdx_ == tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {
440- } else if (coreIdx_ == tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {440+ CalcDgamma(inputOffset, tiling_->colsLastCoreDG, isWithPad);
441- CalcDgamma(inputOffset, tiling_->colsLastCoreDG, isWithPad);441+ }
442- }442+ }
443- }443+ 
444- 444+ __aicore__ inline void ProcessWithLargeRows()
445- __aicore__ inline void ProcessWithLargeRows()445+ {
446- {446+ if (coreIdx_ >= tiling_->usedCoreNumDG) {
447- if (coreIdx_ >= tiling_->usedCoreNumDG) {447+ return;
448- return;448+ }
449- }449+ curCols_ = colsPerUB_;
450- curCols_ = colsPerUB_;450+ if ((coreIdx_ == tiling_->usedCoreNumDG - 1) && !tiling_->isMultiColset) {
451- if ((coreIdx_ == tiling_->usedCoreNumDG - 1) && !tiling_->isMultiColset) {451+ curCols_ = tiling_->colsPerTailCoreDG;
452- curCols_ = tiling_->colsPerTailCoreDG;452+ }
453- }453+ int64_t inputOffset = 0;
454- int64_t inputOffset = 0;454+ int64_t outputOffset = 0;
455- int64_t outputOffset = 0;455+ bool isPowerofTwoRows = (tiling_->tailBlockCountwithPadDG + tiling_->tailBlockCountWithoutPadDG) == 0 ? true :
456- bool isPowerofTwoRows = (tiling_->tailBlockCountwithPadDG + tiling_->tailBlockCountWithoutPadDG) == 0 ? true :456+ false;
457- false;457+ if (isPowerofTwoRows) {
458- if (isPowerofTwoRows) {458+ for (uint32_t curLoop = 0; curLoop < colsUbLoopCount_; curLoop++) {
459- for (uint32_t curLoop = 0; curLoop < colsUbLoopCount_; curLoop++) {459+ outputOffset = curLoop * colsPerUB_ + gmOffset_;
460- outputOffset = curLoop * colsPerUB_ + gmOffset_;460+ HandlingLargeRows2KAlign(outputOffset, curCols_);
461- HandlingLargeRows2KAlign(outputOffset, curCols_);461+ }
462- }462+ outputOffset = tiling_->usedCoreNumDG * colsPerCore_ + coreIdx_ * vlFp32_;
463- outputOffset = tiling_->usedCoreNumDG * colsPerCore_ + coreIdx_ * vlFp32_;463+ if (coreIdx_ < tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {
464- if (coreIdx_ < tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {464+ HandlingLargeRows2KAlign(outputOffset, curCols_);
465- HandlingLargeRows2KAlign(outputOffset, curCols_);465+ } else if (coreIdx_ == tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {
466- } else if (coreIdx_ == tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {466+ HandlingLargeRows2KAlign(outputOffset, tiling_->colsLastCoreDG);
467- HandlingLargeRows2KAlign(outputOffset, tiling_->colsLastCoreDG);467+ }
468- }468+ } else {
469- } else {469+ // 核间累加，2k不对齐场景
470- // 核间累加，2k不对齐场景470+ for (uint32_t curLoop = 0; curLoop < colsUbLoopCount_; curLoop++) {
471- for (uint32_t curLoop = 0; curLoop < colsUbLoopCount_; curLoop++) {471+ uint32_t startOffset = curLoop * colsPerUB_ + gmOffset_;
472- uint32_t startOffset = curLoop * colsPerUB_ + gmOffset_;472+ HandlingLargeRows2KUnAlign(startOffset, curCols_);
473- HandlingLargeRows2KUnAlign(startOffset, curCols_);473+ }
474- }474+ uint32_t startOffset = tiling_->usedCoreNumDG * colsPerCore_ + coreIdx_ * vlFp32_;
475- uint32_t startOffset = tiling_->usedCoreNumDG * colsPerCore_ + coreIdx_ * vlFp32_;475+ if (coreIdx_ < tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {
476- if (coreIdx_ < tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {476+ HandlingLargeRows2KUnAlign(startOffset, curCols_);
477- HandlingLargeRows2KUnAlign(startOffset, curCols_);477+ } else if (coreIdx_ == tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {
478- } else if (coreIdx_ == tiling_->tailCoreNumDG - 1 && tiling_->isMultiColset) {478+ HandlingLargeRows2KUnAlign(startOffset, tiling_->colsLastCoreDG);
479- HandlingLargeRows2KUnAlign(startOffset, tiling_->colsLastCoreDG);479+ }
480- }480+ }
481- }481+ }
482- }482+ 
483- 483+private:
484-private:484+ TQue<QuePosition::VECIN, 1> dyQueue_;
485- TQue<QuePosition::VECIN, 1> dyQueue_;485+ TQue<QuePosition::VECIN, 1> xQueue_;
486- TQue<QuePosition::VECIN, 1> xQueue_;486+ TQue<QuePosition::VECIN, 1> rstdQueue_;
487- TQue<QuePosition::VECIN, 1> rstdQueue_;487+ TQue<QuePosition::VECOUT, 1> dgammaQueue_;
488- TQue<QuePosition::VECOUT, 1> dgammaQueue_;488+ TQue<QuePosition::VECOUT, 1> dgammaQueue1_;
489- TQue<QuePosition::VECOUT, 1> dgammaQueue1_;489+ TQue<QuePosition::VECOUT, 1> binaryAddCacheQueue_;
490- TQue<QuePosition::VECOUT, 1> binaryAddCacheQueue_;490+ 
491- 491+ GlobalTensor<DTYPE_DY> dyGm_;
492- GlobalTensor<DTYPE_DY> dyGm_;492+ GlobalTensor<DTYPE_X> xGm_;
493- GlobalTensor<DTYPE_X> xGm_;493+ GlobalTensor<DTYPE_RSTD> rstdGm_;
494- GlobalTensor<DTYPE_RSTD> rstdGm_;494+ GlobalTensor<float> dgammaGm_;
495- GlobalTensor<float> dgammaGm_;495+ 
496- 496+ uint32_t blockSize_;
497- uint32_t blockSize_;497+ uint32_t vlFp32_;
498- uint32_t vlFp32_;498+ uint32_t coreIdx_;
499- uint32_t coreIdx_;499+ int64_t cols_;
500- int64_t cols_;500+ int64_t rows_;
501- int64_t rows_;501+ uint32_t colsPerUB_;
502- uint32_t colsPerUB_;502+ uint32_t colsPerCore_;
503- uint32_t colsPerCore_;503+ uint32_t rowsPerUB_;
504- uint32_t rowsPerUB_;504+ uint32_t colsUbLoopCount_;
505- uint32_t colsUbLoopCount_;505+ uint32_t colsPerLoopAlign_;
506- uint32_t colsPerLoopAlign_;506+ uint32_t binaryAddK_;
507- uint32_t binaryAddK_;507+ uint32_t tailDataOffset_;
508- uint32_t tailDataOffset_;508+ uint64_t gmOffset_;
509- uint64_t gmOffset_;509+ uint32_t curCols_;
510- uint32_t curCols_;510+ TPipe* Ppipe_ = nullptr;
511- TPipe* Ppipe_ = nullptr;511+ 
512- 512+ const RmsNormGradQuantRegbaseTilingData* tiling_;
513- const RmsNormGradQuantRegbaseTilingData* tiling_;513+};
514-};514+} // namespace RmsNormGradQuant
515-} // namespace RmsNormGradQuant515+#endif // RMS_NORM_GRAD_REGBASE_DGAMMA_H
516-#endif // RMS_NORM_GRAD_REGBASE_DGAMMA_H
Mnorm/rms_norm_grad_quant/op_kernel/arch35/rms_norm_grad_quant_dgamma_big_m.h+491-491
@@ -1,491 +1,491 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/* !11+/* !
12- * \file rms_norm_grad_quant_dgamma_big_m.h12+ * \file rms_norm_grad_quant_dgamma_big_m.h
13- * \brief13+ * \brief
14- */14+ */
15- 15+ 
16-#ifndef RMS_NORM_GRAD_QUANT_DGAMMA_BIG_M_16+#ifndef RMS_NORM_GRAD_QUANT_DGAMMA_BIG_M_
17-#define RMS_NORM_GRAD_QUANT_DGAMMA_BIG_M_17+#define RMS_NORM_GRAD_QUANT_DGAMMA_BIG_M_
18- 18+ 
19-#include "kernel_tiling/kernel_tiling.h"19+#include "kernel_tiling/kernel_tiling.h"
20-#include "kernel_operator.h"20+#include "kernel_operator.h"
21-#include "../inc/kernel_utils.h"21+#include "../inc/kernel_utils.h"
22-#include "rms_norm_grad_quant_dgamma_helper.h"22+#include "rms_norm_grad_quant_dgamma_helper.h"
23-#include "rms_norm_grad_quant_common.h"23+#include "rms_norm_grad_quant_common.h"
24- 24+ 
25-namespace RmsNormGradQuant {25+namespace RmsNormGradQuant {
26- 26+ 
27-template <typename T>27+template <typename T>
28-class RmsNormGradQuantDgammaBigM {28+class RmsNormGradQuantDgammaBigM {
29- static constexpr uint32_t VECTOR_LENGTH = platform::GetVRegSize();29+ static constexpr uint32_t VECTOR_LENGTH = platform::GetVRegSize();
30- static constexpr uint32_t VL_FP32 = VECTOR_LENGTH / sizeof(float);30+ static constexpr uint32_t VL_FP32 = VECTOR_LENGTH / sizeof(float);
31- static constexpr int64_t BLOCK_SIZE = platform::GetUbBlockSize();31+ static constexpr int64_t BLOCK_SIZE = platform::GetUbBlockSize();
32- static constexpr int64_t TRIPLE_BUFFER = 3;32+ static constexpr int64_t TRIPLE_BUFFER = 3;
33- static constexpr int64_t DOUBLE_BUFFER = 2;33+ static constexpr int64_t DOUBLE_BUFFER = 2;
34- 34+ 
35-public:35+public:
36- __aicore__ inline RmsNormGradQuantDgammaBigM(){};36+ __aicore__ inline RmsNormGradQuantDgammaBigM(){};
37- __aicore__ inline void Init(GM_ADDR dy, GM_ADDR x, GM_ADDR rstd, GM_ADDR pdGamma, GM_ADDR workspace,37+ __aicore__ inline void Init(GM_ADDR dy, GM_ADDR x, GM_ADDR rstd, GM_ADDR pdGamma, GM_ADDR workspace,
38- const RmsNormGradQuantRegbaseBigMTilingData* tilingData, TPipe* pipeIn)38+ const RmsNormGradQuantRegbaseBigMTilingData* tilingData, TPipe* pipeIn)
39- {39+ {
40- td_ = tilingData;40+ td_ = tilingData;
41- blockIdx_ = GetBlockIdx();41+ blockIdx_ = GetBlockIdx();
42- if (blockIdx_ >= td_->dgammaUsedCoreNum) {42+ if (blockIdx_ >= td_->dgammaUsedCoreNum) {
43- return;43+ return;
44- }44+ }
45- 45+ 
46- const int64_t startM = (blockIdx_ * td_->dgammaMPerBlock) +46+ const int64_t startM = (blockIdx_ * td_->dgammaMPerBlock) +
47- (blockIdx_ < td_->dgammaMReminder ? blockIdx_ : td_->dgammaMReminder);47+ (blockIdx_ < td_->dgammaMReminder ? blockIdx_ : td_->dgammaMReminder);
48- 48+ 
49- if (blockIdx_ < td_->dgammaMReminder) {49+ if (blockIdx_ < td_->dgammaMReminder) {
50- M = td_->dgammaMToProcessMainBlock;50+ M = td_->dgammaMToProcessMainBlock;
51- Mloop = td_->dgammaMLoopMainBlock;51+ Mloop = td_->dgammaMLoopMainBlock;
52- MTotalLoop = td_->dgammaMTotalLoopMainBlock;52+ MTotalLoop = td_->dgammaMTotalLoopMainBlock;
53- Mtail = td_->dgammaMTailMainBlock;53+ Mtail = td_->dgammaMTailMainBlock;
54- BasicBlockLoop = td_->dgammaBasicBlockLoopMainBlock;54+ BasicBlockLoop = td_->dgammaBasicBlockLoopMainBlock;
55- MainFoldCount = td_->dgammaMainFoldCountMainBlock;55+ MainFoldCount = td_->dgammaMainFoldCountMainBlock;
56- CacheBufferCount = td_->dgammaCacheBufferCountMainBlock;56+ CacheBufferCount = td_->dgammaCacheBufferCountMainBlock;
57- ResultCacheID = td_->dgammaResultCacheIDMainBlock;57+ ResultCacheID = td_->dgammaResultCacheIDMainBlock;
58- } else {58+ } else {
59- M = td_->dgammaMToProcessTailBlock;59+ M = td_->dgammaMToProcessTailBlock;
60- Mloop = td_->dgammaMLoopTailBlock;60+ Mloop = td_->dgammaMLoopTailBlock;
61- MTotalLoop = td_->dgammaMTotalLoopTailBlock;61+ MTotalLoop = td_->dgammaMTotalLoopTailBlock;
62- Mtail = td_->dgammaMTailTailBlock;62+ Mtail = td_->dgammaMTailTailBlock;
63- BasicBlockLoop = td_->dgammaBasicBlockLoopTailBlock;63+ BasicBlockLoop = td_->dgammaBasicBlockLoopTailBlock;
64- MainFoldCount = td_->dgammaMainFoldCountTailBlock;64+ MainFoldCount = td_->dgammaMainFoldCountTailBlock;
65- CacheBufferCount = td_->dgammaCacheBufferCountTailBlock;65+ CacheBufferCount = td_->dgammaCacheBufferCountTailBlock;
66- ResultCacheID = td_->dgammaResultCacheIDTailBlock;66+ ResultCacheID = td_->dgammaResultCacheIDTailBlock;
67- }67+ }
68- 68+ 
69- // Init GM69+ // Init GM
70- int64_t dyOffset = startM * td_->dxTilingData.cols;70+ int64_t dyOffset = startM * td_->dxTilingData.cols;
71- dyInGm_.SetGlobalBuffer((__gm__ T*)dy + dyOffset);71+ dyInGm_.SetGlobalBuffer((__gm__ T*)dy + dyOffset);
72- xInGm_.SetGlobalBuffer((__gm__ T*)x + dyOffset);72+ xInGm_.SetGlobalBuffer((__gm__ T*)x + dyOffset);
73- rstdInGm_.SetGlobalBuffer((__gm__ float*)rstd + startM);73+ rstdInGm_.SetGlobalBuffer((__gm__ float*)rstd + startM);
74- 74+ 
75- int64_t colOffset = blockIdx_ * td_->dxTilingData.cols;75+ int64_t colOffset = blockIdx_ * td_->dxTilingData.cols;
76- dgammaTmpGm_.SetGlobalBuffer((__gm__ float*)workspace + colOffset);76+ dgammaTmpGm_.SetGlobalBuffer((__gm__ float*)workspace + colOffset);
77- dgammaGm_.SetGlobalBuffer((__gm__ float*)pdGamma);77+ dgammaGm_.SetGlobalBuffer((__gm__ float*)pdGamma);
78- 78+ 
79- // Init Pipe79+ // Init Pipe
80- pipe_ = pipeIn;80+ pipe_ = pipeIn;
81- 81+ 
82- int64_t dyBufLen = td_->dgammaMfactorBlockAligned * td_->dgammaNfactorBlockAligned;82+ int64_t dyBufLen = td_->dgammaMfactorBlockAligned * td_->dgammaNfactorBlockAligned;
83- pipe_->InitBuffer(inQueueDy_, TRIPLE_BUFFER, dyBufLen * sizeof(T));83+ pipe_->InitBuffer(inQueueDy_, TRIPLE_BUFFER, dyBufLen * sizeof(T));
84- pipe_->InitBuffer(inQueueX_, TRIPLE_BUFFER, dyBufLen * sizeof(T));84+ pipe_->InitBuffer(inQueueX_, TRIPLE_BUFFER, dyBufLen * sizeof(T));
85- pipe_->InitBuffer(dgammaCalcBuf_, dyBufLen * sizeof(float));85+ pipe_->InitBuffer(dgammaCalcBuf_, dyBufLen * sizeof(float));
86- 86+ 
87- pipe_->InitBuffer(inQueueRstd_, TRIPLE_BUFFER, td_->dgammaMfactorBlockAligned * sizeof(float));87+ pipe_->InitBuffer(inQueueRstd_, TRIPLE_BUFFER, td_->dgammaMfactorBlockAligned * sizeof(float));
88- int64_t nFactorAlignedBufSize = td_->dgammaNfactorBlockAligned * sizeof(float);88+ int64_t nFactorAlignedBufSize = td_->dgammaNfactorBlockAligned * sizeof(float);
89- pipe_->InitBuffer(dgammaTmpOutQueue_, DOUBLE_BUFFER, nFactorAlignedBufSize);89+ pipe_->InitBuffer(dgammaTmpOutQueue_, DOUBLE_BUFFER, nFactorAlignedBufSize);
90- pipe_->InitBuffer(reduceOutTmpBuffer_, nFactorAlignedBufSize);90+ pipe_->InitBuffer(reduceOutTmpBuffer_, nFactorAlignedBufSize);
91- 91+ 
92- int64_t cacheBufSize = CacheBufferCount * nFactorAlignedBufSize;92+ int64_t cacheBufSize = CacheBufferCount * nFactorAlignedBufSize;
93- pipe_->InitBuffer(reduceCacheBuffer_, cacheBufSize);93+ pipe_->InitBuffer(reduceCacheBuffer_, cacheBufSize);
94- }94+ }
95- 95+ 
96- __aicore__ inline void Process()96+ __aicore__ inline void Process()
97- {97+ {
98- // 核内计算98+ // 核内计算
99- if (blockIdx_ < td_->dgammaUsedCoreNum) {99+ if (blockIdx_ < td_->dgammaUsedCoreNum) {
100- ProcessStg0();100+ ProcessStg0();
101- }101+ }
102- 102+ 
103- SyncAll();103+ SyncAll();
104- 104+ 
105- // 0核做核间累加105+ // 0核做核间累加
106- if (blockIdx_ != 0) {106+ if (blockIdx_ != 0) {
107- return;107+ return;
108- }108+ }
109- 109+ 
110- InitBufferStg1();110+ InitBufferStg1();
111- ProcessStg1();111+ ProcessStg1();
112- }112+ }
113- 113+ 
114-private:114+private:
115- // 核内计算115+ // 核内计算
116- __aicore__ inline void ProcessStg0()116+ __aicore__ inline void ProcessStg0()
117- {117+ {
118- dgammaCalcTensor_ = dgammaCalcBuf_.Get<float>();118+ dgammaCalcTensor_ = dgammaCalcBuf_.Get<float>();
119- 119+ 
120- reduceOutTmpTensor_ = reduceOutTmpBuffer_.Get<float>();120+ reduceOutTmpTensor_ = reduceOutTmpBuffer_.Get<float>();
121- reduceCacheTensor_ = reduceCacheBuffer_.Get<float>();121+ reduceCacheTensor_ = reduceCacheBuffer_.Get<float>();
122- 122+ 
123- int64_t totalRounds = td_->dgammaNloop + (td_->dgammaNtail > 0 ? 1 : 0);123+ int64_t totalRounds = td_->dgammaNloop + (td_->dgammaNtail > 0 ? 1 : 0);
124- 124+ 
125- int64_t mfactor = BasicBlockLoop ?125+ int64_t mfactor = BasicBlockLoop ?
126- td_->dgammaMfactorBlockAligned :126+ td_->dgammaMfactorBlockAligned :
127- (M == td_->dgammaMfactorBlockAligned ? td_->dgammaMfactorBlockAligned : Mtail);127+ (M == td_->dgammaMfactorBlockAligned ? td_->dgammaMfactorBlockAligned : Mtail);
128- int64_t loopCnt = BasicBlockLoop ? BasicBlockLoop : 1;128+ int64_t loopCnt = BasicBlockLoop ? BasicBlockLoop : 1;
129- 129+ 
130- for (int64_t round = 0; round < totalRounds; ++round) {130+ for (int64_t round = 0; round < totalRounds; ++round) {
131- int64_t ni = (round < td_->dgammaNloop) ? round : td_->dgammaNloop;131+ int64_t ni = (round < td_->dgammaNloop) ? round : td_->dgammaNloop;
132- int64_t nfactor = (round < td_->dgammaNloop) ? td_->dgammaNfactorBlockAligned : td_->dgammaNtail;132+ int64_t nfactor = (round < td_->dgammaNloop) ? td_->dgammaNfactorBlockAligned : td_->dgammaNtail;
133- 133+ 
134- for (int64_t i = 0; i < loopCnt; ++i) {134+ for (int64_t i = 0; i < loopCnt; ++i) {
135- ProcessMainBlock(ni, i, mfactor, nfactor);135+ ProcessMainBlock(ni, i, mfactor, nfactor);
136- if (BasicBlockLoop != 0 && ((i < MainFoldCount) || (i == MainFoldCount && Mtail > 0))) {136+ if (BasicBlockLoop != 0 && ((i < MainFoldCount) || (i == MainFoldCount && Mtail > 0))) {
137- ProcessFoldBlock(ni, i, (i < MainFoldCount) ? td_->dgammaMfactorBlockAligned : Mtail, nfactor);137+ ProcessFoldBlock(ni, i, (i < MainFoldCount) ? td_->dgammaMfactorBlockAligned : Mtail, nfactor);
138- }138+ }
139- ProcessSummation(ni, i, mfactor, td_->dgammaNfactorBlockAligned);139+ ProcessSummation(ni, i, mfactor, td_->dgammaNfactorBlockAligned);
140- }140+ }
141- 141+ 
142- LocalTensor<float> dgammaTmpOutTensor = dgammaTmpOutQueue_.template AllocTensor<float>();142+ LocalTensor<float> dgammaTmpOutTensor = dgammaTmpOutQueue_.template AllocTensor<float>();
143- 143+ 
144- CopyUB2UB(dgammaTmpOutTensor, reduceCacheTensor_[ResultCacheID * td_->dgammaNfactorBlockAligned], nfactor);144+ CopyUB2UB(dgammaTmpOutTensor, reduceCacheTensor_[ResultCacheID * td_->dgammaNfactorBlockAligned], nfactor);
145- 145+ 
146- dgammaTmpOutQueue_.EnQue(dgammaTmpOutTensor);146+ dgammaTmpOutQueue_.EnQue(dgammaTmpOutTensor);
147- 147+ 
148- int64_t offset = ni * td_->dgammaNfactorBlockAligned;148+ int64_t offset = ni * td_->dgammaNfactorBlockAligned;
149- CopyOutDgammaTmpStg0(offset, nfactor);149+ CopyOutDgammaTmpStg0(offset, nfactor);
150- }150+ }
151- }151+ }
152- 152+ 
153- __aicore__ inline void ProcessMainBlock(const int64_t ni, const int64_t basicBlockIdx, const int64_t mfactor,153+ __aicore__ inline void ProcessMainBlock(const int64_t ni, const int64_t basicBlockIdx, const int64_t mfactor,
154- const int64_t nfactor)154+ const int64_t nfactor)
155- {155+ {
156- int64_t dyOffset = ni * td_->dgammaNfactorBlockAligned +156+ int64_t dyOffset = ni * td_->dgammaNfactorBlockAligned +
157- basicBlockIdx * td_->dgammaMfactorBlockAligned * td_->dxTilingData.cols;157+ basicBlockIdx * td_->dgammaMfactorBlockAligned * td_->dxTilingData.cols;
158- 158+ 
159- CopyInDyAndX(dyOffset, mfactor, nfactor, td_->dgammaNfactorBlockAligned, td_->dxTilingData.cols);159+ CopyInDyAndX(dyOffset, mfactor, nfactor, td_->dgammaNfactorBlockAligned, td_->dxTilingData.cols);
160- 160+ 
161- int64_t rstdOffset = basicBlockIdx * td_->dgammaMfactorBlockAligned;161+ int64_t rstdOffset = basicBlockIdx * td_->dgammaMfactorBlockAligned;
162- 162+ 
163- CopyInRstd(rstdOffset, mfactor);163+ CopyInRstd(rstdOffset, mfactor);
164- 164+ 
165- LocalTensor<T> dyTensor = inQueueDy_.template DeQue<T>();165+ LocalTensor<T> dyTensor = inQueueDy_.template DeQue<T>();
166- LocalTensor<T> xTensor = inQueueX_.template DeQue<T>();166+ LocalTensor<T> xTensor = inQueueX_.template DeQue<T>();
167- LocalTensor<float> rstdTensor = inQueueRstd_.template DeQue<float>();167+ LocalTensor<float> rstdTensor = inQueueRstd_.template DeQue<float>();
168- ComputeDgammaTmpMain(dgammaCalcTensor_, dyTensor, xTensor, rstdTensor, mfactor, td_->dgammaNfactorBlockAligned);168+ ComputeDgammaTmpMain(dgammaCalcTensor_, dyTensor, xTensor, rstdTensor, mfactor, td_->dgammaNfactorBlockAligned);
169- inQueueDy_.FreeTensor(dyTensor);169+ inQueueDy_.FreeTensor(dyTensor);
170- inQueueX_.FreeTensor(xTensor);170+ inQueueX_.FreeTensor(xTensor);
171- inQueueRstd_.FreeTensor(rstdTensor);171+ inQueueRstd_.FreeTensor(rstdTensor);
172- }172+ }
173- 173+ 
174- __aicore__ inline void ProcessFoldBlock(const int64_t ni, const int64_t basicBlockIdx, const int64_t mfactor,174+ __aicore__ inline void ProcessFoldBlock(const int64_t ni, const int64_t basicBlockIdx, const int64_t mfactor,
175- const int64_t nfactor)175+ const int64_t nfactor)
176- {176+ {
177- int64_t dyOffset = ni * td_->dgammaNfactorBlockAligned +177+ int64_t dyOffset = ni * td_->dgammaNfactorBlockAligned +
178- (basicBlockIdx + BasicBlockLoop) * td_->dgammaMfactorBlockAligned * td_->dxTilingData.cols;178+ (basicBlockIdx + BasicBlockLoop) * td_->dgammaMfactorBlockAligned * td_->dxTilingData.cols;
179- 179+ 
180- CopyInDyAndX(dyOffset, mfactor, nfactor, td_->dgammaNfactorBlockAligned, td_->dxTilingData.cols);180+ CopyInDyAndX(dyOffset, mfactor, nfactor, td_->dgammaNfactorBlockAligned, td_->dxTilingData.cols);
181- int64_t rstdOffset = (basicBlockIdx + BasicBlockLoop) * td_->dgammaMfactorBlockAligned;181+ int64_t rstdOffset = (basicBlockIdx + BasicBlockLoop) * td_->dgammaMfactorBlockAligned;
182- CopyInRstd(rstdOffset, mfactor);182+ CopyInRstd(rstdOffset, mfactor);
183- 183+ 
184- LocalTensor<T> dyTensor = inQueueDy_.template DeQue<T>();184+ LocalTensor<T> dyTensor = inQueueDy_.template DeQue<T>();
185- LocalTensor<T> xTensor = inQueueX_.template DeQue<T>();185+ LocalTensor<T> xTensor = inQueueX_.template DeQue<T>();
186- LocalTensor<float> rstdTensor = inQueueRstd_.template DeQue<float>();186+ LocalTensor<float> rstdTensor = inQueueRstd_.template DeQue<float>();
187- ComputeDgammaTmpFold(dgammaCalcTensor_, dyTensor, xTensor, rstdTensor, mfactor, td_->dgammaNfactorBlockAligned);187+ ComputeDgammaTmpFold(dgammaCalcTensor_, dyTensor, xTensor, rstdTensor, mfactor, td_->dgammaNfactorBlockAligned);
188- inQueueDy_.FreeTensor(dyTensor);188+ inQueueDy_.FreeTensor(dyTensor);
189- inQueueX_.FreeTensor(xTensor);189+ inQueueX_.FreeTensor(xTensor);
190- inQueueRstd_.FreeTensor(rstdTensor);190+ inQueueRstd_.FreeTensor(rstdTensor);
191- }191+ }
192- 192+ 
193- __aicore__ inline void ProcessSummation(const int64_t ni, const int64_t basicBlockIdx, const int64_t mfactor,193+ __aicore__ inline void ProcessSummation(const int64_t ni, const int64_t basicBlockIdx, const int64_t mfactor,
194- const int64_t nfactor)194+ const int64_t nfactor)
195- {195+ {
196- int64_t cacheID = GetCacheID(basicBlockIdx);196+ int64_t cacheID = GetCacheID(basicBlockIdx);
197- uint32_t srcShape[2] = {static_cast<uint32_t>(mfactor), static_cast<uint32_t>(nfactor)};197+ uint32_t srcShape[2] = {static_cast<uint32_t>(mfactor), static_cast<uint32_t>(nfactor)};
198- 198+ 
199- AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(reduceOutTmpTensor_, dgammaCalcTensor_, srcShape,199+ AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(reduceOutTmpTensor_, dgammaCalcTensor_, srcShape,
200- false);200+ false);
201- 201+ 
202- UpdateCache(reduceCacheTensor_, reduceOutTmpTensor_, cacheID, td_->dgammaNfactorBlockAligned, nfactor);202+ UpdateCache(reduceCacheTensor_, reduceOutTmpTensor_, cacheID, td_->dgammaNfactorBlockAligned, nfactor);
203- }203+ }
204- 204+ 
205- __aicore__ inline void ComputeDgammaTmpMain(const LocalTensor<float>& dstTensor, const LocalTensor<T>& dyTensor,205+ __aicore__ inline void ComputeDgammaTmpMain(const LocalTensor<float>& dstTensor, const LocalTensor<T>& dyTensor,
206- const LocalTensor<T>& xTensor, const LocalTensor<float>& rstdTensor,206+ const LocalTensor<T>& xTensor, const LocalTensor<float>& rstdTensor,
207- const int64_t rowSize, const int64_t colSize)207+ const int64_t rowSize, const int64_t colSize)
208- {208+ {
209- uint16_t outerLoopTimes = static_cast<uint16_t>(rowSize);209+ uint16_t outerLoopTimes = static_cast<uint16_t>(rowSize);
210- uint16_t innerLoopTimes = ops::CeilDiv(colSize, static_cast<int64_t>(VL_FP32));210+ uint16_t innerLoopTimes = ops::CeilDiv(colSize, static_cast<int64_t>(VL_FP32));
211- uint32_t outerStride = td_->dgammaNfactorBlockAligned;211+ uint32_t outerStride = td_->dgammaNfactorBlockAligned;
212- uint32_t innerStride = VL_FP32;212+ uint32_t innerStride = VL_FP32;
213- 213+ 
214- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();214+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
215- __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr();215+ __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr();
216- __local_mem__ T* dy = (__local_mem__ T*)dyTensor.GetPhyAddr();216+ __ubuf__ T* dy = (__ubuf__ T*)dyTensor.GetPhyAddr();
217- __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr();217+ __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr();
218- 218+ 
219- __VEC_SCOPE__219+ __VEC_SCOPE__
220- {220+ {
221- for (uint16_t i = 0; i < outerLoopTimes; ++i) {221+ for (uint16_t i = 0; i < outerLoopTimes; ++i) {
222- uint32_t count = static_cast<uint32_t>(colSize);222+ uint32_t count = static_cast<uint32_t>(colSize);
223- AscendC::MicroAPI::RegTensor<float> rstdReg;223+ AscendC::MicroAPI::RegTensor<float> rstdReg;
224- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(224+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(
225- rstdReg, (__local_mem__ float*)rstd + static_cast<uint32_t>(i));225+ rstdReg, (__ubuf__ float*)rstd + static_cast<uint32_t>(i));
226- 226+ 
227- AscendC::MicroAPI::RegTensor<float> xReg;227+ AscendC::MicroAPI::RegTensor<float> xReg;
228- AscendC::MicroAPI::RegTensor<float> dyReg;228+ AscendC::MicroAPI::RegTensor<float> dyReg;
229- AscendC::MicroAPI::MaskReg pMask;229+ AscendC::MicroAPI::MaskReg pMask;
230- for (uint16_t j = 0; j < innerLoopTimes; ++j) {230+ for (uint16_t j = 0; j < innerLoopTimes; ++j) {
231- pMask = AscendC::MicroAPI::UpdateMask<float>(count);231+ pMask = AscendC::MicroAPI::UpdateMask<float>(count);
232- uint32_t offset = i * outerStride + j * innerStride;232+ uint32_t offset = i * outerStride + j * innerStride;
233- LoadOneTensor<T>(xReg, x, pMask, offset);233+ LoadOneTensor<T>(xReg, x, pMask, offset);
234- Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask);234+ Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask);
235- LoadOneTensor<T>(dyReg, dy, pMask, offset);235+ LoadOneTensor<T>(dyReg, dy, pMask, offset);
236- Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);236+ Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);
237- DataCopy((__local_mem__ float*)dst + offset, xReg, pMask);237+ StoreAlign((__ubuf__ float*)dst + offset, xReg, pMask);
238- }238+ }
239- }239+ }
240- }240+ }
241- }241+ }
242- 242+ 
243- __aicore__ inline void ComputeDgammaTmpFold(const LocalTensor<float>& dstTensor, const LocalTensor<T>& dyTensor,243+ __aicore__ inline void ComputeDgammaTmpFold(const LocalTensor<float>& dstTensor, const LocalTensor<T>& dyTensor,
244- const LocalTensor<T>& xTensor, const LocalTensor<float>& rstdTensor,244+ const LocalTensor<T>& xTensor, const LocalTensor<float>& rstdTensor,
245- const int64_t rowSize, const int64_t colSize)245+ const int64_t rowSize, const int64_t colSize)
246- {246+ {
247- uint16_t outerLoopTimes = static_cast<uint16_t>(rowSize);247+ uint16_t outerLoopTimes = static_cast<uint16_t>(rowSize);
248- uint16_t innerLoopTimes = ops::CeilDiv(colSize, static_cast<int64_t>(VL_FP32));248+ uint16_t innerLoopTimes = ops::CeilDiv(colSize, static_cast<int64_t>(VL_FP32));
249- uint32_t outerStride = td_->dgammaNfactorBlockAligned;249+ uint32_t outerStride = td_->dgammaNfactorBlockAligned;
250- uint32_t innerStride = VL_FP32;250+ uint32_t innerStride = VL_FP32;
251- 251+ 
252- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();252+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
253- __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr();253+ __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr();
254- __local_mem__ T* dy = (__local_mem__ T*)dyTensor.GetPhyAddr();254+ __ubuf__ T* dy = (__ubuf__ T*)dyTensor.GetPhyAddr();
255- __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr();255+ __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr();
256- 256+ 
257- __VEC_SCOPE__257+ __VEC_SCOPE__
258- {258+ {
259- for (uint16_t i = 0; i < outerLoopTimes; ++i) {259+ for (uint16_t i = 0; i < outerLoopTimes; ++i) {
260- uint32_t count = static_cast<uint32_t>(colSize);260+ uint32_t count = static_cast<uint32_t>(colSize);
261- AscendC::MicroAPI::RegTensor<float> rstdReg;261+ AscendC::MicroAPI::RegTensor<float> rstdReg;
262- DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(262+ LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(
263- rstdReg, (__local_mem__ float*)rstd + static_cast<uint32_t>(i));263+ rstdReg, (__ubuf__ float*)rstd + static_cast<uint32_t>(i));
264- 264+ 
265- AscendC::MicroAPI::RegTensor<float> xReg;265+ AscendC::MicroAPI::RegTensor<float> xReg;
266- AscendC::MicroAPI::RegTensor<float> dyReg;266+ AscendC::MicroAPI::RegTensor<float> dyReg;
267- AscendC::MicroAPI::MaskReg pMask;267+ AscendC::MicroAPI::MaskReg pMask;
268- for (uint16_t j = 0; j < innerLoopTimes; ++j) {268+ for (uint16_t j = 0; j < innerLoopTimes; ++j) {
269- pMask = AscendC::MicroAPI::UpdateMask<float>(count);269+ pMask = AscendC::MicroAPI::UpdateMask<float>(count);
270- uint32_t offset = i * outerStride + j * innerStride;270+ uint32_t offset = i * outerStride + j * innerStride;
271- LoadOneTensor<T>(xReg, x, pMask, offset);271+ LoadOneTensor<T>(xReg, x, pMask, offset);
272- Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask);272+ Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask);
273- LoadOneTensor<T>(dyReg, dy, pMask, offset);273+ LoadOneTensor<T>(dyReg, dy, pMask, offset);
274- Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);274+ Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);
275- LoadOneTensor<float>(dyReg, dst, pMask, offset);275+ LoadOneTensor<float>(dyReg, dst, pMask, offset);
276- Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);276+ Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask);
277- DataCopy((__local_mem__ float*)dst + offset, xReg, pMask);277+ StoreAlign((__ubuf__ float*)dst + offset, xReg, pMask);
278- }278+ }
279- }279+ }
280- }280+ }
281- }281+ }
282- 282+ 
283- template <typename T1>283+ template <typename T1>
284- __aicore__ inline void LoadOneTensor(MicroAPI::RegTensor<float>& dst, const __local_mem__ void* input,284+ __aicore__ inline void LoadOneTensor(MicroAPI::RegTensor<float>& dst, const __ubuf__ void* input,
285- MicroAPI::MaskReg& preg, uint32_t offset)285+ MicroAPI::MaskReg& preg, uint32_t offset)
286- {286+ {
287- if constexpr (!IsSameType<T1, float>::value) {287+ if constexpr (!IsSameType<T1, float>::value) {
288- MicroAPI::RegTensor<T1> xFp16;288+ MicroAPI::RegTensor<T1> xFp16;
289- DataCopy<T1, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, (__local_mem__ T1*)(input) + offset);289+ LoadAlign<T1, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, (__ubuf__ T1*)(input) + offset);
290- Cast<float, T1, castTraitB162B32>(dst, xFp16, preg);290+ Cast<float, T1, castTraitB162B32>(dst, xFp16, preg);
291- } else {291+ } else {
292- DataCopy(dst, (__local_mem__ float*)(input) + offset);292+ LoadAlign(dst, (__ubuf__ float*)(input) + offset);
293- }293+ }
294- }294+ }
295- 295+ 
296- __aicore__ inline void CopyOutDgammaTmpStg0(int64_t offset, const int64_t curAInnerLen)296+ __aicore__ inline void CopyOutDgammaTmpStg0(int64_t offset, const int64_t curAInnerLen)
297- {297+ {
298- DataCopyExtParams copyOutParams;298+ DataCopyExtParams copyOutParams;
299- copyOutParams.blockCount = 1;299+ copyOutParams.blockCount = 1;
300- copyOutParams.blockLen = curAInnerLen * sizeof(float);300+ copyOutParams.blockLen = curAInnerLen * sizeof(float);
301- copyOutParams.srcStride = 0;301+ copyOutParams.srcStride = 0;
302- copyOutParams.dstStride = 0;302+ copyOutParams.dstStride = 0;
303- 303+ 
304- LocalTensor<float> dgammaTmp = dgammaTmpOutQueue_.template DeQue<float>();304+ LocalTensor<float> dgammaTmp = dgammaTmpOutQueue_.template DeQue<float>();
305- DataCopyPad<float, PaddingMode::Normal>(dgammaTmpGm_[offset], dgammaTmp, copyOutParams);305+ DataCopyPad<float, PaddingMode::Normal>(dgammaTmpGm_[offset], dgammaTmp, copyOutParams);
306- dgammaTmpOutQueue_.FreeTensor(dgammaTmp);306+ dgammaTmpOutQueue_.FreeTensor(dgammaTmp);
307- }307+ }
308- 308+ 
309- __aicore__ inline void CopyInDyAndX(const int64_t gmOffset, const int64_t nburst, const int64_t burstLen,309+ __aicore__ inline void CopyInDyAndX(const int64_t gmOffset, const int64_t nburst, const int64_t burstLen,
310- const int64_t dstStride, const int64_t srcStride)310+ const int64_t dstStride, const int64_t srcStride)
311- {311+ {
312- DataCopyExtParams params;312+ DataCopyExtParams params;
313- params.blockCount = nburst;313+ params.blockCount = nburst;
314- params.blockLen = burstLen * sizeof(T);314+ params.blockLen = burstLen * sizeof(T);
315- params.srcStride = srcStride * sizeof(T) - params.blockLen;315+ params.srcStride = srcStride * sizeof(T) - params.blockLen;
316- params.dstStride = (dstStride - burstLen) * sizeof(T) / BLOCK_SIZE;316+ params.dstStride = (dstStride - burstLen) * sizeof(T) / BLOCK_SIZE;
317- 317+ 
318- DataCopyPadExtParams<T> padParams;318+ DataCopyPadExtParams<T> padParams;
319- padParams.isPad = false;319+ padParams.isPad = false;
320- 320+ 
321- LocalTensor<T> dyTensor = inQueueDy_.AllocTensor<T>();321+ LocalTensor<T> dyTensor = inQueueDy_.AllocTensor<T>();
322- DataCopyPad<T, PaddingMode::Normal>(dyTensor, dyInGm_[gmOffset], params, padParams);322+ DataCopyPad<T, PaddingMode::Normal>(dyTensor, dyInGm_[gmOffset], params, padParams);
323- inQueueDy_.EnQue(dyTensor);323+ inQueueDy_.EnQue(dyTensor);
324- 324+ 
325- LocalTensor<T> xTensor = inQueueX_.AllocTensor<T>();325+ LocalTensor<T> xTensor = inQueueX_.AllocTensor<T>();
326- DataCopyPad<T, PaddingMode::Normal>(xTensor, xInGm_[gmOffset], params, padParams);326+ DataCopyPad<T, PaddingMode::Normal>(xTensor, xInGm_[gmOffset], params, padParams);
327- inQueueX_.EnQue(xTensor);327+ inQueueX_.EnQue(xTensor);
328- }328+ }
329- 329+ 
330- __aicore__ inline void CopyInRstd(const int64_t gmOffset, const int64_t burstLen)330+ __aicore__ inline void CopyInRstd(const int64_t gmOffset, const int64_t burstLen)
331- {331+ {
332- DataCopyExtParams params;332+ DataCopyExtParams params;
333- params.blockCount = 1;333+ params.blockCount = 1;
334- params.blockLen = burstLen * sizeof(float);334+ params.blockLen = burstLen * sizeof(float);
335- 335+ 
336- DataCopyPadExtParams<float> padParams;336+ DataCopyPadExtParams<float> padParams;
337- padParams.isPad = false;337+ padParams.isPad = false;
338- 338+ 
339- LocalTensor<float> rstdTensor = inQueueRstd_.template AllocTensor<float>();339+ LocalTensor<float> rstdTensor = inQueueRstd_.template AllocTensor<float>();
340- DataCopyPad<float, PaddingMode::Normal>(rstdTensor, rstdInGm_[gmOffset], params, padParams);340+ DataCopyPad<float, PaddingMode::Normal>(rstdTensor, rstdInGm_[gmOffset], params, padParams);
341- inQueueRstd_.EnQue(rstdTensor);341+ inQueueRstd_.EnQue(rstdTensor);
342- }342+ }
343- 343+ 
344- __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,344+ __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,
345- const int64_t cacheID, const int64_t stride, const int64_t count)345+ const int64_t cacheID, const int64_t stride, const int64_t count)
346- {346+ {
347- uint16_t outerLoopTimes = ops::CeilDiv(static_cast<uint32_t>(count), VL_FP32);347+ uint16_t outerLoopTimes = ops::CeilDiv(static_cast<uint32_t>(count), VL_FP32);
348- uint16_t innerLoopTimes = cacheID;348+ uint16_t innerLoopTimes = cacheID;
349- uint32_t outerLoopStride = VL_FP32;349+ uint32_t outerLoopStride = VL_FP32;
350- uint32_t innerLoopStride = stride;350+ uint32_t innerLoopStride = stride;
351- __VEC_SCOPE__351+ __VEC_SCOPE__
352- {352+ {
353- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();353+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
354- __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride;354+ __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride;
355- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();355+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
356- uint32_t sreg = static_cast<uint32_t>(count);356+ uint32_t sreg = static_cast<uint32_t>(count);
357- AscendC::MicroAPI::RegTensor<float> aReg, bReg;357+ AscendC::MicroAPI::RegTensor<float> aReg, bReg;
358- AscendC::MicroAPI::MaskReg pMask;358+ AscendC::MicroAPI::MaskReg pMask;
359- for (uint16_t i = 0; i < outerLoopTimes; ++i) {359+ for (uint16_t i = 0; i < outerLoopTimes; ++i) {
360- pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);360+ pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
361- DataCopy(aReg, (__local_mem__ float*)src + static_cast<uint32_t>(i * outerLoopStride));361+ LoadAlign(aReg, (__ubuf__ float*)src + static_cast<uint32_t>(i * outerLoopStride));
362- for (uint16_t j = 0; j < innerLoopTimes; ++j) {362+ for (uint16_t j = 0; j < innerLoopTimes; ++j) {
363- DataCopy(bReg, (__local_mem__ float*)dst +363+ LoadAlign(bReg,
364- static_cast<uint32_t>(i * outerLoopStride + j * innerLoopStride));364+ (__ubuf__ float*)dst + static_cast<uint32_t>(i * outerLoopStride + j * innerLoopStride));
365- Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);365+ Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
366- }366+ }
367- DataCopy((__local_mem__ float*)cah + static_cast<uint32_t>(i * outerLoopStride), aReg, pMask);367+ StoreAlign((__ubuf__ float*)cah + static_cast<uint32_t>(i * outerLoopStride), aReg, pMask);
368- }368+ }
369- }369+ }
370- }370+ }
371- 371+ 
372- __aicore__ inline void CopyUB2UB(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,372+ __aicore__ inline void CopyUB2UB(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor,
373- const int64_t count)373+ const int64_t count)
374- {374+ {
375- DataCopy(dstTensor, srcTensor,375+ DataCopy(dstTensor, srcTensor,
376- ops::Aligned(static_cast<int64_t>(count), static_cast<int64_t>(BLOCK_SIZE / sizeof(float))));376+ ops::Aligned(static_cast<int64_t>(count), static_cast<int64_t>(BLOCK_SIZE / sizeof(float))));
377- }377+ }
378- 378+ 
379- // 核间累加379+ // 核间累加
380- __aicore__ inline void InitBufferStg1()380+ __aicore__ inline void InitBufferStg1()
381- {381+ {
382- pipe_->Reset();382+ pipe_->Reset();
383- int64_t aSizeIn = td_->dgammaAInnerAlignedStg1 * td_->dgammaUsedCoreNum * sizeof(float);383+ int64_t aSizeIn = td_->dgammaAInnerAlignedStg1 * td_->dgammaUsedCoreNum * sizeof(float);
384- pipe_->InitBuffer(dgammaTmpInQue_, DOUBLE_BUFFER, aSizeIn);384+ pipe_->InitBuffer(dgammaTmpInQue_, DOUBLE_BUFFER, aSizeIn);
385- 385+ 
386- pipe_->InitBuffer(dgammaOutQue_, DOUBLE_BUFFER, td_->dgammaAInnerAlignedStg1 * sizeof(float));386+ pipe_->InitBuffer(dgammaOutQue_, DOUBLE_BUFFER, td_->dgammaAInnerAlignedStg1 * sizeof(float));
387- }387+ }
388- 388+ 
389- __aicore__ inline void ProcessStg1()389+ __aicore__ inline void ProcessStg1()
390- {390+ {
391- for (int64_t i = 0; i < td_->dgammaAOuterStg1; i++) {391+ for (int64_t i = 0; i < td_->dgammaAOuterStg1; i++) {
392- uint16_t curAInnerLen = i != (td_->dgammaAOuterStg1 - 1) ? td_->dgammaAInnerAlignedStg1 :392+ uint16_t curAInnerLen = i != (td_->dgammaAOuterStg1 - 1) ? td_->dgammaAInnerAlignedStg1 :
393- td_->dgammaATailStg1;393+ td_->dgammaATailStg1;
394- 394+ 
395- int64_t offset = i * td_->dgammaAInnerAlignedStg1;395+ int64_t offset = i * td_->dgammaAInnerAlignedStg1;
396- CopyInDgammaTmp(offset, curAInnerLen);396+ CopyInDgammaTmp(offset, curAInnerLen);
397- 397+ 
398- LocalTensor<float> dgammaTmpIn = dgammaTmpInQue_.template DeQue<float>();398+ LocalTensor<float> dgammaTmpIn = dgammaTmpInQue_.template DeQue<float>();
399- LocalTensor<float> dgammaOut = dgammaOutQue_.AllocTensor<float>();399+ LocalTensor<float> dgammaOut = dgammaOutQue_.AllocTensor<float>();
400- 400+ 
401- uint32_t srcShape[2] = {static_cast<uint32_t>(td_->dgammaUsedCoreNum),401+ uint32_t srcShape[2] = {static_cast<uint32_t>(td_->dgammaUsedCoreNum),
402- static_cast<uint32_t>(td_->dgammaAInnerAlignedStg1)};402+ static_cast<uint32_t>(td_->dgammaAInnerAlignedStg1)};
403- 403+ 
404- ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(dgammaOut, dgammaTmpIn, srcShape, false);404+ ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(dgammaOut, dgammaTmpIn, srcShape, false);
405- dgammaTmpInQue_.FreeTensor(dgammaTmpIn);405+ dgammaTmpInQue_.FreeTensor(dgammaTmpIn);
406- dgammaOutQue_.EnQue(dgammaOut);406+ dgammaOutQue_.EnQue(dgammaOut);
407- CopyOutDgamma(offset, curAInnerLen);407+ CopyOutDgamma(offset, curAInnerLen);
408- }408+ }
409- }409+ }
410- 410+ 
411- __aicore__ inline void CopyInDgammaTmp(int64_t offset, int64_t curALen)411+ __aicore__ inline void CopyInDgammaTmp(int64_t offset, int64_t curALen)
412- {412+ {
413- DataCopyPadExtParams<float> dataCopyPadExtParams;413+ DataCopyPadExtParams<float> dataCopyPadExtParams;
414- dataCopyPadExtParams.isPad = false;414+ dataCopyPadExtParams.isPad = false;
415- dataCopyPadExtParams.leftPadding = 0;415+ dataCopyPadExtParams.leftPadding = 0;
416- dataCopyPadExtParams.rightPadding = 0;416+ dataCopyPadExtParams.rightPadding = 0;
417- dataCopyPadExtParams.paddingValue = 0;417+ dataCopyPadExtParams.paddingValue = 0;
418- 418+ 
419- DataCopyExtParams copyInParams;419+ DataCopyExtParams copyInParams;
420- copyInParams.blockCount = td_->dgammaUsedCoreNum;420+ copyInParams.blockCount = td_->dgammaUsedCoreNum;
421- copyInParams.blockLen = curALen * sizeof(float);421+ copyInParams.blockLen = curALen * sizeof(float);
422- copyInParams.srcStride = (td_->dxTilingData.cols - curALen) * sizeof(float);422+ copyInParams.srcStride = (td_->dxTilingData.cols - curALen) * sizeof(float);
423- copyInParams.dstStride = (td_->dgammaAInnerAlignedStg1 - curALen) * sizeof(float) / BLOCK_SIZE;423+ copyInParams.dstStride = (td_->dgammaAInnerAlignedStg1 - curALen) * sizeof(float) / BLOCK_SIZE;
424- 424+ 
425- LocalTensor<float> dgammaTmpTensor = dgammaTmpInQue_.AllocTensor<float>();425+ LocalTensor<float> dgammaTmpTensor = dgammaTmpInQue_.AllocTensor<float>();
426- DataCopyPad<float, PaddingMode::Normal>(dgammaTmpTensor, dgammaTmpGm_[offset], copyInParams,426+ DataCopyPad<float, PaddingMode::Normal>(dgammaTmpTensor, dgammaTmpGm_[offset], copyInParams,
427- dataCopyPadExtParams);427+ dataCopyPadExtParams);
428- dgammaTmpInQue_.EnQue(dgammaTmpTensor);428+ dgammaTmpInQue_.EnQue(dgammaTmpTensor);
429- }429+ }
430- 430+ 
431- __aicore__ inline void CopyOutDgamma(int64_t offset, uint32_t curAInnerLen)431+ __aicore__ inline void CopyOutDgamma(int64_t offset, uint32_t curAInnerLen)
432- {432+ {
433- DataCopyExtParams copyOutParams;433+ DataCopyExtParams copyOutParams;
434- copyOutParams.blockCount = 1;434+ copyOutParams.blockCount = 1;
435- copyOutParams.blockLen = curAInnerLen * sizeof(float);435+ copyOutParams.blockLen = curAInnerLen * sizeof(float);
436- copyOutParams.srcStride = 0;436+ copyOutParams.srcStride = 0;
437- copyOutParams.dstStride = 0;437+ copyOutParams.dstStride = 0;
438- 438+ 
439- LocalTensor<float> dgamma = dgammaOutQue_.DeQue<float>();439+ LocalTensor<float> dgamma = dgammaOutQue_.DeQue<float>();
440- DataCopyPad<float, PaddingMode::Normal>(dgammaGm_[offset], dgamma, copyOutParams);440+ DataCopyPad<float, PaddingMode::Normal>(dgammaGm_[offset], dgamma, copyOutParams);
441- dgammaOutQue_.FreeTensor(dgamma);441+ dgammaOutQue_.FreeTensor(dgamma);
442- }442+ }
443- 443+ 
444-private:444+private:
445- const RmsNormGradQuantRegbaseBigMTilingData* td_;445+ const RmsNormGradQuantRegbaseBigMTilingData* td_;
446- TPipe* pipe_;446+ TPipe* pipe_;
447- 447+ 
448- int64_t blockIdx_ = 0;448+ int64_t blockIdx_ = 0;
449- 449+ 
450- int64_t M = 0;450+ int64_t M = 0;
451- 451+ 
452- int64_t Mloop = 0;452+ int64_t Mloop = 0;
453- int64_t Mtail = 0;453+ int64_t Mtail = 0;
454- int64_t MTotalLoop = 0;454+ int64_t MTotalLoop = 0;
455- 455+ 
456- int64_t BasicBlockLoop = 0;456+ int64_t BasicBlockLoop = 0;
457- int64_t MainFoldCount = 0;457+ int64_t MainFoldCount = 0;
458- int64_t CacheBufferCount = 0;458+ int64_t CacheBufferCount = 0;
459- int64_t ResultCacheID = 0;459+ int64_t ResultCacheID = 0;
460- 460+ 
461- // Global Tensor461+ // Global Tensor
462- GlobalTensor<T> dyInGm_;462+ GlobalTensor<T> dyInGm_;
463- GlobalTensor<T> xInGm_;463+ GlobalTensor<T> xInGm_;
464- GlobalTensor<float> rstdInGm_;464+ GlobalTensor<float> rstdInGm_;
465- GlobalTensor<float> dgammaTmpGm_;465+ GlobalTensor<float> dgammaTmpGm_;
466- GlobalTensor<float> dgammaGm_;466+ GlobalTensor<float> dgammaGm_;
467- 467+ 
468- // Local Tensor468+ // Local Tensor
469- LocalTensor<float> dgammaCalcTensor_;469+ LocalTensor<float> dgammaCalcTensor_;
470- LocalTensor<float> reduceOutTmpTensor_;470+ LocalTensor<float> reduceOutTmpTensor_;
471- LocalTensor<float> reduceCacheTensor_;471+ LocalTensor<float> reduceCacheTensor_;
472- 472+ 
473- LocalTensor<float> dyMain_;473+ LocalTensor<float> dyMain_;
474- LocalTensor<float> xMain_;474+ LocalTensor<float> xMain_;
475- LocalTensor<float> rstd_;475+ LocalTensor<float> rstd_;
476- 476+ 
477- // TQue477+ // TQue
478- TQue<QuePosition::VECIN, 1> inQueueDy_;478+ TQue<QuePosition::VECIN, 1> inQueueDy_;
479- TQue<QuePosition::VECIN, 1> inQueueX_;479+ TQue<QuePosition::VECIN, 1> inQueueX_;
480- TQue<QuePosition::VECIN, 1> inQueueRstd_;480+ TQue<QuePosition::VECIN, 1> inQueueRstd_;
481- TQue<QuePosition::VECOUT, 1> dgammaTmpOutQueue_;481+ TQue<QuePosition::VECOUT, 1> dgammaTmpOutQueue_;
482- TQue<QuePosition::VECIN, 1> dgammaTmpInQue_;482+ TQue<QuePosition::VECIN, 1> dgammaTmpInQue_;
483- TQue<QuePosition::VECOUT, 1> dgammaOutQue_;483+ TQue<QuePosition::VECOUT, 1> dgammaOutQue_;
484- 484+ 
485- TBuf<TPosition::VECCALC> dgammaCalcBuf_;485+ TBuf<TPosition::VECCALC> dgammaCalcBuf_;
486- TBuf<TPosition::VECCALC> reduceCacheBuffer_;486+ TBuf<TPosition::VECCALC> reduceCacheBuffer_;
487- TBuf<TPosition::VECCALC> reduceOutTmpBuffer_;487+ TBuf<TPosition::VECCALC> reduceOutTmpBuffer_;
488-}; // RmsNormGradQuantDgammaBigM488+}; // RmsNormGradQuantDgammaBigM
489- 489+ 
490-} // namespace RmsNormGradQuant490+} // namespace RmsNormGradQuant
491-#endif // RMS_NORM_GRAD_QUANT_DGAMMA_BIG_M_491+#endif // RMS_NORM_GRAD_QUANT_DGAMMA_BIG_M_
Mnorm/rms_norm_grad_quant/op_kernel/arch35/rms_norm_grad_quant_dgamma_helper.h+245-246
@@ -1,246 +1,245 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/*!11+/*!
12- * \file rms_norm_grad_quant_dgamma_helper.h12+ * \file rms_norm_grad_quant_dgamma_helper.h
13- * \brief RmsNormGrad regbase dgamma helper file13+ * \brief RmsNormGrad regbase dgamma helper file
14- */14+ */
15-#ifndef RMS_NORM_GRAD_DGAMMA_QUANT_HELPER_H15+#ifndef RMS_NORM_GRAD_DGAMMA_QUANT_HELPER_H
16-#define RMS_NORM_GRAD_DGAMMA_QUANT_HELPER_H16+#define RMS_NORM_GRAD_DGAMMA_QUANT_HELPER_H
17- 17+ 
18-#include "kernel_tiling/kernel_tiling.h"18+#include "kernel_tiling/kernel_tiling.h"
19-#include "kernel_operator.h"19+#include "kernel_operator.h"
20-#include "rms_norm_grad_quant_common.h"20+#include "rms_norm_grad_quant_common.h"
21-namespace RmsNormGradQuant {21+namespace RmsNormGradQuant {
22-using namespace AscendC;22+using namespace AscendC;
23-using AscendC::MicroAPI::CreateMask;23+using AscendC::MicroAPI::CreateMask;
24-using AscendC::MicroAPI::LoadDist;24+using AscendC::MicroAPI::LoadDist;
25-using AscendC::MicroAPI::MaskPattern;25+using AscendC::MicroAPI::MaskPattern;
26-using AscendC::MicroAPI::MaskReg;26+using AscendC::MicroAPI::MaskReg;
27-using AscendC::MicroAPI::MemType;27+using AscendC::MicroAPI::MemType;
28-using AscendC::MicroAPI::RegTensor;28+using AscendC::MicroAPI::RegTensor;
29-using AscendC::MicroAPI::StoreDist;29+using AscendC::MicroAPI::StoreDist;
30-using AscendC::MicroAPI::UpdateMask;30+using AscendC::MicroAPI::UpdateMask;
31- 31+ 
32-constexpr int32_t REDUCEBY8ELENUM = 16;32+constexpr int32_t REDUCEBY8ELENUM = 16;
33-constexpr int32_t REDUCEBY4ELENUM = 8;33+constexpr int32_t REDUCEBY4ELENUM = 8;
34-constexpr int32_t REDUCEBY2ELENUM = 4;34+constexpr int32_t REDUCEBY2ELENUM = 4;
35-constexpr int32_t REDUCEBY1ELENUM = 2;35+constexpr int32_t REDUCEBY1ELENUM = 2;
36-constexpr int32_t COMPRESSBY8ELENUM = 8;36+constexpr int32_t COMPRESSBY8ELENUM = 8;
37-constexpr int32_t RESERVESIZE = 32;37+constexpr int32_t RESERVESIZE = 32;
38- 38+ 
39-__aicore__ inline int64_t CEIL_DIV(int64_t x, int64_t y) { return (y > 0) ? (x + y - 1) / y : 0; }39+__aicore__ inline int64_t CEIL_DIV(int64_t x, int64_t y) { return (y > 0) ? (x + y - 1) / y : 0; }
40- 40+ 
41-__aicore__ inline uint32_t BLOCK_ALIGN(uint32_t x, uint32_t blockSize)41+__aicore__ inline uint32_t BLOCK_ALIGN(uint32_t x, uint32_t blockSize)
42-{42+{
43- return (blockSize > 0) ? (x + blockSize - 1) / blockSize * blockSize : 0;43+ return (blockSize > 0) ? (x + blockSize - 1) / blockSize * blockSize : 0;
44-}44+}
45- 45+ 
46-template <typename DY_TYPE, typename X_TYPE, typename RSTD_TYPE>46+template <typename DY_TYPE, typename X_TYPE, typename RSTD_TYPE>
47-__aicore__ inline void CalcMulRes(__local_mem__ DY_TYPE* dyAddr, __local_mem__ X_TYPE* xAddr,47+__aicore__ inline void CalcMulRes(__ubuf__ DY_TYPE* dyAddr, __ubuf__ X_TYPE* xAddr, __ubuf__ RSTD_TYPE* rstdAddr,
48- __local_mem__ RSTD_TYPE* rstdAddr, __local_mem__ float* dgammaOutAddr, MaskReg& preg,48+ __ubuf__ float* dgammaOutAddr, MaskReg& preg, uint32_t offset0, uint32_t k)
49- uint32_t offset0, uint32_t k)49+{
50-{50+ RegTensor<float> xFp32, rstdFp32, dyFp32, temp_res, mul_res;
51- RegTensor<float> xFp32, rstdFp32, dyFp32, temp_res, mul_res;51+ 
52- 52+ if constexpr (IsSameType<DY_TYPE, float>::value) {
53- if constexpr (IsSameType<DY_TYPE, float>::value) {53+ LoadAlign<DY_TYPE, LoadDist::DIST_NORM>(dyFp32, (__ubuf__ float*)(dyAddr + offset0));
54- DataCopy<DY_TYPE, LoadDist::DIST_NORM>(dyFp32, (__local_mem__ float*)(dyAddr + offset0));54+ } else {
55- } else {55+ RegTensor<DY_TYPE> dstRegB16;
56- RegTensor<DY_TYPE> dstRegB16;56+ LoadAlign<DY_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__ubuf__ DY_TYPE*)(dyAddr + offset0));
57- DataCopy<DY_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__local_mem__ DY_TYPE*)(dyAddr + offset0));57+ Cast<float, DY_TYPE, castTraitB162B32>(dyFp32, dstRegB16, preg);
58- Cast<float, DY_TYPE, castTraitB162B32>(dyFp32, dstRegB16, preg);58+ }
59- }59+ 
60- 60+ if constexpr (IsSameType<X_TYPE, float>::value) {
61- if constexpr (IsSameType<X_TYPE, float>::value) {61+ LoadAlign<X_TYPE, LoadDist::DIST_NORM>(xFp32, (__ubuf__ float*)(xAddr + offset0));
62- DataCopy<X_TYPE, LoadDist::DIST_NORM>(xFp32, (__local_mem__ float*)(xAddr + offset0));62+ } else {
63- } else {63+ RegTensor<X_TYPE> dstRegB16;
64- RegTensor<X_TYPE> dstRegB16;64+ LoadAlign<X_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__ubuf__ X_TYPE*)(xAddr + offset0));
65- DataCopy<X_TYPE, LoadDist::DIST_UNPACK_B16>(dstRegB16, (__local_mem__ X_TYPE*)(xAddr + offset0));65+ Cast<float, X_TYPE, castTraitB162B32>(xFp32, dstRegB16, preg);
66- Cast<float, X_TYPE, castTraitB162B32>(xFp32, dstRegB16, preg);66+ }
67- }67+ 
68- 68+ LoadAlign<RSTD_TYPE, LoadDist::DIST_BRC_B32>(rstdFp32, ((__ubuf__ float*)rstdAddr + k));
69- DataCopy<RSTD_TYPE, LoadDist::DIST_BRC_B32>(rstdFp32, ((__local_mem__ float*)rstdAddr + k));69+ 
70- 70+ Mul(temp_res, xFp32, rstdFp32, preg);
71- Mul(temp_res, xFp32, rstdFp32, preg);71+ Mul(mul_res, dyFp32, temp_res, preg);
72- Mul(mul_res, dyFp32, temp_res, preg);72+ 
73- 73+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dgammaOutAddr + offset0), mul_res, preg);
74- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dgammaOutAddr + offset0), mul_res, preg);74+}
75-}75+ 
76- 76+__aicore__ inline void reduceSumCompressedBy8(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset,
77-__aicore__ inline void reduceSumCompressedBy8(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset,77+ uint32_t ub_offset)
78- uint32_t ub_offset)78+{
79-{79+ RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,
80- RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,80+ temp_reg3_1, temp_reg4_0, temp_reg4_1, temp_reg5_0, temp_reg5_1, temp_reg6_0, temp_reg6_1, temp_reg7_0,
81- temp_reg3_1, temp_reg4_0, temp_reg4_1, temp_reg5_0, temp_reg5_1, temp_reg6_0, temp_reg6_1, temp_reg7_0,81+ temp_reg7_1;
82- temp_reg7_1;82+ __ubuf__ float* currentAddr = dyAddr + REDUCEBY8ELENUM * ub_offset;
83- __local_mem__ float* currentAddr = dyAddr + REDUCEBY8ELENUM * ub_offset;83+ 
84- 84+ //
85- //85+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(currentAddr));
86- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(currentAddr));86+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(currentAddr + offset));
87- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(currentAddr + offset));87+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
88- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);88+ 
89- 89+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_0, (__ubuf__ float*)(currentAddr + 2 * offset));
90- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_0, (__local_mem__ float*)(currentAddr + 2 * offset));90+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_1, (__ubuf__ float*)(currentAddr + 3 * offset));
91- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_1, (__local_mem__ float*)(currentAddr + 3 * offset));91+ AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);
92- AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);92+ 
93- 93+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_0, (__ubuf__ float*)(currentAddr + 4 * offset));
94- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_0, (__local_mem__ float*)(currentAddr + 4 * offset));94+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_1, (__ubuf__ float*)(currentAddr + 5 * offset));
95- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_1, (__local_mem__ float*)(currentAddr + 5 * offset));95+ AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);
96- AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);96+ 
97- 97+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_0, (__ubuf__ float*)(currentAddr + 6 * offset));
98- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_0, (__local_mem__ float*)(currentAddr + 6 * offset));98+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_1, (__ubuf__ float*)(currentAddr + 7 * offset));
99- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_1, (__local_mem__ float*)(currentAddr + 7 * offset));99+ AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);
100- AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);100+ 
101- 101+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg4_0, (__ubuf__ float*)(currentAddr + 8 * offset));
102- DataCopy<float, LoadDist::DIST_NORM>(temp_reg4_0, (__local_mem__ float*)(currentAddr + 8 * offset));102+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg4_1, (__ubuf__ float*)(currentAddr + 9 * offset));
103- DataCopy<float, LoadDist::DIST_NORM>(temp_reg4_1, (__local_mem__ float*)(currentAddr + 9 * offset));103+ AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg4_1, preg);
104- AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg4_1, preg);104+ 
105- 105+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg5_0, (__ubuf__ float*)(currentAddr + 10 * offset));
106- DataCopy<float, LoadDist::DIST_NORM>(temp_reg5_0, (__local_mem__ float*)(currentAddr + 10 * offset));106+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg5_1, (__ubuf__ float*)(currentAddr + 11 * offset));
107- DataCopy<float, LoadDist::DIST_NORM>(temp_reg5_1, (__local_mem__ float*)(currentAddr + 11 * offset));107+ AscendC::MicroAPI::Add(temp_reg5_0, temp_reg5_0, temp_reg5_1, preg);
108- AscendC::MicroAPI::Add(temp_reg5_0, temp_reg5_0, temp_reg5_1, preg);108+ 
109- 109+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg6_0, (__ubuf__ float*)(currentAddr + 12 * offset));
110- DataCopy<float, LoadDist::DIST_NORM>(temp_reg6_0, (__local_mem__ float*)(currentAddr + 12 * offset));110+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg6_1, (__ubuf__ float*)(currentAddr + 13 * offset));
111- DataCopy<float, LoadDist::DIST_NORM>(temp_reg6_1, (__local_mem__ float*)(currentAddr + 13 * offset));111+ AscendC::MicroAPI::Add(temp_reg6_0, temp_reg6_0, temp_reg6_1, preg);
112- AscendC::MicroAPI::Add(temp_reg6_0, temp_reg6_0, temp_reg6_1, preg);112+ 
113- 113+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg7_0, (__ubuf__ float*)(currentAddr + 14 * offset));
114- DataCopy<float, LoadDist::DIST_NORM>(temp_reg7_0, (__local_mem__ float*)(currentAddr + 14 * offset));114+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg7_1, (__ubuf__ float*)(currentAddr + 15 * offset));
115- DataCopy<float, LoadDist::DIST_NORM>(temp_reg7_1, (__local_mem__ float*)(currentAddr + 15 * offset));115+ AscendC::MicroAPI::Add(temp_reg7_0, temp_reg7_0, temp_reg7_1, preg);
116- AscendC::MicroAPI::Add(temp_reg7_0, temp_reg7_0, temp_reg7_1, preg);116+ 
117- 117+ //
118- //118+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);
119- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);119+ AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg3_0, preg);
120- AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg3_0, preg);120+ AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg5_0, preg);
121- AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg5_0, preg);121+ AscendC::MicroAPI::Add(temp_reg6_0, temp_reg6_0, temp_reg7_0, preg);
122- AscendC::MicroAPI::Add(temp_reg6_0, temp_reg6_0, temp_reg7_0, preg);122+ 
123- 123+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg2_0, preg);
124- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg2_0, preg);124+ AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg6_0, preg);
125- AscendC::MicroAPI::Add(temp_reg4_0, temp_reg4_0, temp_reg6_0, preg);125+ 
126- 126+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg4_0, preg);
127- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg4_0, preg);127+ 
128- 128+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);
129- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);129+}
130-}130+ 
131- 131+__aicore__ inline void reduceSumCompressedBy4(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset,
132-__aicore__ inline void reduceSumCompressedBy4(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset,132+ uint32_t ub_offset)
133- uint32_t ub_offset)133+{
134-{134+ RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,
135- RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1, temp_reg2_0, temp_reg2_1, temp_reg3_0,135+ temp_reg3_1;
136- temp_reg3_1;136+ __ubuf__ float* currentAddr = dyAddr + COMPRESSBY8ELENUM * ub_offset;
137- __local_mem__ float* currentAddr = dyAddr + COMPRESSBY8ELENUM * ub_offset;137+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(currentAddr));
138- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(currentAddr));138+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(currentAddr + offset));
139- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(currentAddr + offset));139+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
140- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);140+ 
141- 141+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_0, (__ubuf__ float*)(currentAddr + 2 * offset));
142- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_0, (__local_mem__ float*)(currentAddr + 2 * offset));142+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_1, (__ubuf__ float*)(currentAddr + 3 * offset));
143- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_1, (__local_mem__ float*)(currentAddr + 3 * offset));143+ AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);
144- AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);144+ 
145- 145+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_0, (__ubuf__ float*)(currentAddr + 4 * offset));
146- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_0, (__local_mem__ float*)(currentAddr + 4 * offset));146+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg2_1, (__ubuf__ float*)(currentAddr + 5 * offset));
147- DataCopy<float, LoadDist::DIST_NORM>(temp_reg2_1, (__local_mem__ float*)(currentAddr + 5 * offset));147+ AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);
148- AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg2_1, preg);148+ 
149- 149+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_0, (__ubuf__ float*)(currentAddr + 6 * offset));
150- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_0, (__local_mem__ float*)(currentAddr + 6 * offset));150+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg3_1, (__ubuf__ float*)(currentAddr + 7 * offset));
151- DataCopy<float, LoadDist::DIST_NORM>(temp_reg3_1, (__local_mem__ float*)(currentAddr + 7 * offset));151+ AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);
152- AscendC::MicroAPI::Add(temp_reg3_0, temp_reg3_0, temp_reg3_1, preg);152+ 
153- 153+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);
154- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);154+ AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg3_0, preg);
155- AscendC::MicroAPI::Add(temp_reg2_0, temp_reg2_0, temp_reg3_0, preg);155+ 
156- 156+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg2_0, preg);
157- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg2_0, preg);157+ 
158- 158+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);
159- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);159+}
160-}160+ 
161- 161+__aicore__ inline void reduceSumCompressedBy2(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset,
162-__aicore__ inline void reduceSumCompressedBy2(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset,162+ uint32_t ub_offset)
163- uint32_t ub_offset)163+{
164-{164+ RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1;
165- RegTensor<float> temp_reg0_0, temp_reg0_1, temp_reg1_0, temp_reg1_1;165+ 
166- 166+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(dyAddr));
167- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(dyAddr));167+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(dyAddr + offset));
168- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(dyAddr + offset));168+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
169- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);169+ 
170- 170+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_0, (__ubuf__ float*)(dyAddr + 2 * offset));
171- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_0, (__local_mem__ float*)(dyAddr + 2 * offset));171+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg1_1, (__ubuf__ float*)(dyAddr + 3 * offset));
172- DataCopy<float, LoadDist::DIST_NORM>(temp_reg1_1, (__local_mem__ float*)(dyAddr + 3 * offset));172+ AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);
173- AscendC::MicroAPI::Add(temp_reg1_0, temp_reg1_0, temp_reg1_1, preg);173+ 
174- 174+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);
175- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg1_0, preg);175+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);
176- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr + ub_offset), temp_reg0_0, preg);176+}
177-}177+ 
178- 178+__aicore__ inline void reduceSumCompressedBy1(__ubuf__ float* dyAddr, MaskReg& preg, uint32_t offset)
179-__aicore__ inline void reduceSumCompressedBy1(__local_mem__ float* dyAddr, MaskReg& preg, uint32_t offset)179+{
180-{180+ RegTensor<float> temp_reg0_0, temp_reg0_1;
181- RegTensor<float> temp_reg0_0, temp_reg0_1;181+ 
182- 182+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(dyAddr));
183- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(dyAddr));183+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(dyAddr + offset));
184- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(dyAddr + offset));184+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
185- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);185+ 
186- 186+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(dyAddr), temp_reg0_0, preg);
187- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(dyAddr), temp_reg0_0, preg);187+}
188-}188+ 
189- 189+__aicore__ inline void reduceSumCompressedBy8WithOutPad(__ubuf__ float* src1Addr, __ubuf__ float* src2Addr,
190-__aicore__ inline void reduceSumCompressedBy8WithOutPad(__local_mem__ float* src1Addr, __local_mem__ float* src2Addr,190+ MaskReg& preg, uint32_t ub_offset, uint32_t vlFp32)
191- MaskReg& preg, uint32_t ub_offset, uint32_t vlFp32)191+{
192-{192+ for (uint16_t i = 0; i < 8; i++) {
193- for (uint16_t i = 0; i < 8; i++) {193+ RegTensor<float> temp_reg0_0, temp_reg0_1;
194- RegTensor<float> temp_reg0_0, temp_reg0_1;194+ uint32_t tempOffset = i * vlFp32;
195- uint32_t tempOffset = i * vlFp32;195+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(src1Addr + ub_offset + tempOffset));
196- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(src1Addr + ub_offset + tempOffset));196+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(src2Addr + ub_offset + tempOffset));
197- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(src2Addr + ub_offset + tempOffset));197+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
198- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);198+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(src1Addr + ub_offset + tempOffset), temp_reg0_0,
199- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(src1Addr + ub_offset + tempOffset),199+ preg);
200- temp_reg0_0, preg);200+ }
201- }201+}
202-}202+ 
203- 203+__aicore__ inline void reduceSumCompressedBy8WithPad(__ubuf__ float* src1Addr, __ubuf__ float* src2Addr, MaskReg& preg,
204-__aicore__ inline void reduceSumCompressedBy8WithPad(__local_mem__ float* src1Addr, __local_mem__ float* src2Addr,204+ uint32_t ub_offset, uint32_t rowsBoundLine, uint32_t vlFp32,
205- MaskReg& preg, uint32_t ub_offset, uint32_t rowsBoundLine,205+ uint32_t tailDataOffset)
206- uint32_t vlFp32, uint32_t tailDataOffset)206+{
207-{207+ for (uint16_t i = 0; i < 8; i++) {
208- for (uint16_t i = 0; i < 8; i++) {208+ RegTensor<float> temp_reg0_0, temp_reg0_1;
209- RegTensor<float> temp_reg0_0, temp_reg0_1;209+ uint32_t temp_off_set_0 = ub_offset + i * vlFp32;
210- uint32_t temp_off_set_0 = ub_offset + i * vlFp32;210+ uint32_t temp_off_set_1 = tailDataOffset + temp_off_set_0 < rowsBoundLine ? tailDataOffset + temp_off_set_0 :
211- uint32_t temp_off_set_1 = tailDataOffset + temp_off_set_0 < rowsBoundLine ? tailDataOffset + temp_off_set_0 :211+ rowsBoundLine;
212- rowsBoundLine;212+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_0, (__ubuf__ float*)(src1Addr + temp_off_set_0));
213- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_0, (__local_mem__ float*)(src1Addr + temp_off_set_0));213+ LoadAlign<float, LoadDist::DIST_NORM>(temp_reg0_1, (__ubuf__ float*)(src2Addr + temp_off_set_1));
214- DataCopy<float, LoadDist::DIST_NORM>(temp_reg0_1, (__local_mem__ float*)(src2Addr + temp_off_set_1));214+ 
215- 215+ AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);
216- AscendC::MicroAPI::Add(temp_reg0_0, temp_reg0_0, temp_reg0_1, preg);216+ StoreAlign<float, StoreDist::DIST_NORM_B32>((__ubuf__ float*)(src1Addr + temp_off_set_0), temp_reg0_0, preg);
217- DataCopy<float, StoreDist::DIST_NORM_B32>((__local_mem__ float*)(src1Addr + temp_off_set_0), temp_reg0_0, preg);217+ }
218- }218+}
219-}219+ 
220- 220+__aicore__ inline void UpdateCache(const AscendC::LocalTensor<float>& dstTensor, __ubuf__ float* srcAddr,
221-__aicore__ inline void UpdateCache(const AscendC::LocalTensor<float>& dstTensor, __local_mem__ float* srcAddr,221+ const int64_t cacheID, const int64_t count)
222- const int64_t cacheID, const int64_t count)222+{
223-{223+ // UpdateCache
224- // UpdateCache224+ uint16_t innerLoopTimes = cacheID;
225- uint16_t innerLoopTimes = cacheID;225+ uint32_t innerLoopStride = count;
226- uint32_t innerLoopStride = count;226+ __VEC_SCOPE__
227- __VEC_SCOPE__227+ {
228- {228+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
229- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();229+ __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * count;
230- __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * count;230+ uint32_t sreg = static_cast<uint32_t>(count);
231- uint32_t sreg = static_cast<uint32_t>(count);231+ AscendC::MicroAPI::RegTensor<float> aReg, bReg;
232- AscendC::MicroAPI::RegTensor<float> aReg, bReg;232+ AscendC::MicroAPI::MaskReg pMask;
233- AscendC::MicroAPI::MaskReg pMask;233+ pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
234- pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);234+ LoadAlign(aReg, (__ubuf__ float*)srcAddr);
235- DataCopy(aReg, (__local_mem__ float*)srcAddr);235+ for (uint16_t j = 0; j < innerLoopTimes; ++j) {
236- for (uint16_t j = 0; j < innerLoopTimes; ++j) {236+ LoadAlign(bReg, (__ubuf__ float*)dst + static_cast<uint32_t>(j * innerLoopStride));
237- DataCopy(bReg, (__local_mem__ float*)dst + static_cast<uint32_t>(j * innerLoopStride));237+ Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
238- Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);238+ }
239- }239+ StoreAlign((__ubuf__ float*)cah, aReg, pMask);
240- DataCopy((__local_mem__ float*)cah, aReg, pMask);240+ }
241- }241+}
242-}242+ 
243- 243+__aicore__ inline int64_t GetCacheID(const int64_t idx) { return ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1; }
244-__aicore__ inline int64_t GetCacheID(const int64_t idx) { return ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1; }244+} // namespace RmsNormGradQuant
245-} // namespace RmsNormGradQuant245+#endif // RMS_NORM_GRAD_REGBASE_DGAMMA_H
246-#endif // RMS_NORM_GRAD_REGBASE_DGAMMA_H
Mnorm/rms_norm_grad_quant/op_kernel/arch35/rms_norm_grad_quant_dx_full_load.h+456-456
@@ -1,456 +1,456 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/*!11+/*!
12- * \file rms_norm_grad_quant_regbase_dx_full_load.h12+ * \file rms_norm_grad_quant_regbase_dx_full_load.h
13- * \brief RmsNormGradQuant Regbase DX Full Load kernel File13+ * \brief RmsNormGradQuant Regbase DX Full Load kernel File
14- */14+ */
15- 15+ 
16-#ifndef RMS_NORM_GRAD_Quant_DX_FULL_LOAD_H16+#ifndef RMS_NORM_GRAD_Quant_DX_FULL_LOAD_H
17-#define RMS_NORM_GRAD_Quant_DX_FULL_LOAD_H17+#define RMS_NORM_GRAD_Quant_DX_FULL_LOAD_H
18- 18+ 
19-#include "kernel_tiling/kernel_tiling.h"19+#include "kernel_tiling/kernel_tiling.h"
20-#include "kernel_operator.h"20+#include "kernel_operator.h"
21-#include "rms_norm_grad_quant_common.h"21+#include "rms_norm_grad_quant_common.h"
22- 22+ 
23-namespace RmsNormGradQuant {23+namespace RmsNormGradQuant {
24-using namespace AscendC;24+using namespace AscendC;
25-template <typename T_DY, typename T_X, typename T_GAMMA, typename T_DX, typename T_DGAMMA, typename T_SCALES_X,25+template <typename T_DY, typename T_X, typename T_GAMMA, typename T_DX, typename T_DGAMMA, typename T_SCALES_X,
26- typename T_OFFSET_X, bool HAS_OFFSET_X, bool DIV_MODE>26+ typename T_OFFSET_X, bool HAS_OFFSET_X, bool DIV_MODE>
27-class RegbaseDxFullLoad {27+class RegbaseDxFullLoad {
28-public:28+public:
29- __aicore__ inline RegbaseDxFullLoad(TPipe* pipe, const RmsNormGradQuantRegbaseDxTilingData* tilingData)29+ __aicore__ inline RegbaseDxFullLoad(TPipe* pipe, const RmsNormGradQuantRegbaseDxTilingData* tilingData)
30- : Ppipe_(pipe), tiling_(tilingData)30+ : Ppipe_(pipe), tiling_(tilingData)
31- {}31+ {}
32- 32+ 
33- __aicore__ inline void Init(__gm__ uint8_t* dy, __gm__ uint8_t* x, __gm__ uint8_t* rstd, __gm__ uint8_t* gamma,33+ __aicore__ inline void Init(__gm__ uint8_t* dy, __gm__ uint8_t* x, __gm__ uint8_t* rstd, __gm__ uint8_t* gamma,
34- __gm__ uint8_t* scales_x, __gm__ uint8_t* offset_x, __gm__ uint8_t* dx,34+ __gm__ uint8_t* scales_x, __gm__ uint8_t* offset_x, __gm__ uint8_t* dx,
35- __gm__ uint8_t* dgamma)35+ __gm__ uint8_t* dgamma)
36- {36+ {
37-#if (__NPU_ARCH__ == 3510)37+#if (__NPU_ARCH__ == 3510)
38- if constexpr (IsSameType<T_DX, hifloat8_t>::value) {38+ if constexpr (IsSameType<T_DX, hifloat8_t>::value) {
39- AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);39+ AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
40- }40+ }
41-#endif41+#endif
42- usedCoreNum_ = tiling_->usedCoreNumDx;42+ usedCoreNum_ = tiling_->usedCoreNumDx;
43- uint32_t coreIdx = GetBlockIdx();43+ uint32_t coreIdx = GetBlockIdx();
44- if (coreIdx >= usedCoreNum_) {44+ if (coreIdx >= usedCoreNum_) {
45- return;45+ return;
46- }46+ }
47- rows_ = tiling_->rows;47+ rows_ = tiling_->rows;
48- cols_ = tiling_->cols;48+ cols_ = tiling_->cols;
49- blockFactor_ = tiling_->blockFactorDx;49+ blockFactor_ = tiling_->blockFactorDx;
50- 50+ 
51- colsAlignBlock_ = IsSameType<T_X, float>::value ? AlignUp(cols_, FLOAT_NUM_BLOCK) :51+ colsAlignBlock_ = IsSameType<T_X, float>::value ? AlignUp(cols_, FLOAT_NUM_BLOCK) :
52- AlignUp(cols_, HALF_NUM_BLOCK);52+ AlignUp(cols_, HALF_NUM_BLOCK);
53- if constexpr (IsSameType<T_DX, hifloat8_t>::value || IsSameType<T_DX, int8_t>::value) {53+ if constexpr (IsSameType<T_DX, hifloat8_t>::value || IsSameType<T_DX, int8_t>::value) {
54- colsAlignHiFP8_ = AlignUp(cols_, HIFP8_NUM_BLOCK);54+ colsAlignHiFP8_ = AlignUp(cols_, HIFP8_NUM_BLOCK);
55- }55+ }
56- colsAlign2VL_ = AlignUp(cols_, FLOAT_NUM_2VL);56+ colsAlign2VL_ = AlignUp(cols_, FLOAT_NUM_2VL);
57- 57+ 
58- ubFactor_ = UB_FACTOR_DX_FULL_LOAD;58+ ubFactor_ = UB_FACTOR_DX_FULL_LOAD;
59- ubFactorD_ = colsAlign2VL_;59+ ubFactorD_ = colsAlign2VL_;
60- ubFactorN_ = ubFactor_ / ubFactorD_;60+ ubFactorN_ = ubFactor_ / ubFactorD_;
61- avgFactor1_ = 1.0f / cols_;61+ avgFactor1_ = 1.0f / cols_;
62- 62+ 
63- dyGm_.SetGlobalBuffer((__gm__ T_DY*)dy + coreIdx * blockFactor_ * cols_);63+ dyGm_.SetGlobalBuffer((__gm__ T_DY*)dy + coreIdx * blockFactor_ * cols_);
64- xGm_.SetGlobalBuffer((__gm__ T_X*)x + coreIdx * blockFactor_ * cols_);64+ xGm_.SetGlobalBuffer((__gm__ T_X*)x + coreIdx * blockFactor_ * cols_);
65- rstdGm_.SetGlobalBuffer((__gm__ float*)rstd + coreIdx * blockFactor_);65+ rstdGm_.SetGlobalBuffer((__gm__ float*)rstd + coreIdx * blockFactor_);
66- gammaGm_.SetGlobalBuffer((__gm__ T_GAMMA*)gamma);66+ gammaGm_.SetGlobalBuffer((__gm__ T_GAMMA*)gamma);
67- dxGm_.SetGlobalBuffer((__gm__ T_DX*)dx + coreIdx * blockFactor_ * cols_);67+ dxGm_.SetGlobalBuffer((__gm__ T_DX*)dx + coreIdx * blockFactor_ * cols_);
68- 68+ 
69- Ppipe_->InitBuffer(inQueueDy_, DB_NUM, ubFactor_ * sizeof(float));69+ Ppipe_->InitBuffer(inQueueDy_, DB_NUM, ubFactor_ * sizeof(float));
70- Ppipe_->InitBuffer(inQueueX_, DB_NUM, ubFactor_ * sizeof(float));70+ Ppipe_->InitBuffer(inQueueX_, DB_NUM, ubFactor_ * sizeof(float));
71- Ppipe_->InitBuffer(inQueueRstd_, DB_NUM, AlignUp(ubFactorN_, V_LENGTH) * sizeof(float));71+ Ppipe_->InitBuffer(inQueueRstd_, DB_NUM, AlignUp(ubFactorN_, V_LENGTH) * sizeof(float));
72- Ppipe_->InitBuffer(outQueueDx_, DB_NUM, ubFactor_ * sizeof(float));72+ Ppipe_->InitBuffer(outQueueDx_, DB_NUM, ubFactor_ * sizeof(float));
73- Ppipe_->InitBuffer(inQueueGamma_, 1, ubFactor_ * sizeof(float));73+ Ppipe_->InitBuffer(inQueueGamma_, 1, ubFactor_ * sizeof(float));
74- Ppipe_->InitBuffer(reduceBuf_, ubFactorN_ * colsAlign2VL_ * sizeof(float));74+ Ppipe_->InitBuffer(reduceBuf_, ubFactorN_ * colsAlign2VL_ * sizeof(float));
75- Ppipe_->InitBuffer(tmpSumBuf_, AlignUp(ubFactorN_, V_LENGTH) * sizeof(float));75+ Ppipe_->InitBuffer(tmpSumBuf_, AlignUp(ubFactorN_, V_LENGTH) * sizeof(float));
76- scalesXGm_.SetGlobalBuffer((__gm__ T_SCALES_X*)scales_x);76+ scalesXGm_.SetGlobalBuffer((__gm__ T_SCALES_X*)scales_x);
77- Ppipe_->InitBuffer(inQueueScalesX_, 1, sizeof(T_SCALES_X));77+ Ppipe_->InitBuffer(inQueueScalesX_, 1, sizeof(T_SCALES_X));
78- if constexpr (HAS_OFFSET_X) {78+ if constexpr (HAS_OFFSET_X) {
79- offsetXGm_.SetGlobalBuffer((__gm__ T_OFFSET_X*)offset_x);79+ offsetXGm_.SetGlobalBuffer((__gm__ T_OFFSET_X*)offset_x);
80- Ppipe_->InitBuffer(inQueueOffsetX_, 1, sizeof(T_OFFSET_X));80+ Ppipe_->InitBuffer(inQueueOffsetX_, 1, sizeof(T_OFFSET_X));
81- }81+ }
82- }82+ }
83- __aicore__ inline void Process()83+ __aicore__ inline void Process()
84- {84+ {
85- uint32_t coreIdx = GetBlockIdx();85+ uint32_t coreIdx = GetBlockIdx();
86- if (coreIdx >= usedCoreNum_) {86+ if (coreIdx >= usedCoreNum_) {
87- return;87+ return;
88- }88+ }
89- // copyInScalesX89+ // copyInScalesX
90- CopyInScalesX();90+ CopyInScalesX();
91- if constexpr (HAS_OFFSET_X) {91+ if constexpr (HAS_OFFSET_X) {
92- CopyInOffsetX();92+ CopyInOffsetX();
93- }93+ }
94- int64_t blockTail = rows_ - (usedCoreNum_ - 1) * blockFactor_;94+ int64_t blockTail = rows_ - (usedCoreNum_ - 1) * blockFactor_;
95- int64_t calcRowNum = coreIdx == usedCoreNum_ - 1 ? blockTail : blockFactor_;95+ int64_t calcRowNum = coreIdx == usedCoreNum_ - 1 ? blockTail : blockFactor_;
96- int64_t calcRowNumRemain = calcRowNum;96+ int64_t calcRowNumRemain = calcRowNum;
97- for (int64_t rowIdx = 0; rowIdx < calcRowNum; rowIdx += ubFactorN_) {97+ for (int64_t rowIdx = 0; rowIdx < calcRowNum; rowIdx += ubFactorN_) {
98- int64_t calcRowNumSub = Min(ubFactorN_, calcRowNumRemain);98+ int64_t calcRowNumSub = Min(ubFactorN_, calcRowNumRemain);
99- SubProcess(rowIdx, calcRowNumSub);99+ SubProcess(rowIdx, calcRowNumSub);
100- calcRowNumRemain -= ubFactorN_;100+ calcRowNumRemain -= ubFactorN_;
101- }101+ }
102- if (calcRowNum > 0) {102+ if (calcRowNum > 0) {
103- inQueueGamma_.FreeTensor(gammaLocal_);103+ inQueueGamma_.FreeTensor(gammaLocal_);
104- }104+ }
105- inQueueScalesX_.FreeTensor(scalesXLocal_);105+ inQueueScalesX_.FreeTensor(scalesXLocal_);
106- if constexpr (HAS_OFFSET_X) {106+ if constexpr (HAS_OFFSET_X) {
107- inQueueOffsetX_.FreeTensor(offsetXLocal_);107+ inQueueOffsetX_.FreeTensor(offsetXLocal_);
108- }108+ }
109- }109+ }
110- 110+ 
111- __aicore__ inline void SubProcess(int64_t rowIdx, int64_t calcRowNumSub)111+ __aicore__ inline void SubProcess(int64_t rowIdx, int64_t calcRowNumSub)
112- {112+ {
113- if (rowIdx == 0) {113+ if (rowIdx == 0) {
114- CopyInGamma();114+ CopyInGamma();
115- }115+ }
116- CopyInDy(rowIdx, calcRowNumSub);116+ CopyInDy(rowIdx, calcRowNumSub);
117- LocalTensor<float> dyLocal = inQueueDy_.DeQue<float>();117+ LocalTensor<float> dyLocal = inQueueDy_.DeQue<float>();
118- CopyInX(rowIdx, calcRowNumSub);118+ CopyInX(rowIdx, calcRowNumSub);
119- LocalTensor<float> xLocal = inQueueX_.DeQue<float>();119+ LocalTensor<float> xLocal = inQueueX_.DeQue<float>();
120- CopyInRstd(rowIdx, calcRowNumSub);120+ CopyInRstd(rowIdx, calcRowNumSub);
121- LocalTensor<float> rstdLocal = inQueueRstd_.DeQue<float>();121+ LocalTensor<float> rstdLocal = inQueueRstd_.DeQue<float>();
122- LocalTensor<T_GAMMA> gammaLocal = gammaLocal_;122+ LocalTensor<T_GAMMA> gammaLocal = gammaLocal_;
123- LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();123+ LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();
124- 124+ 
125- LocalTensor<float> reduceLocal = reduceBuf_.Get<float>();125+ LocalTensor<float> reduceLocal = reduceBuf_.Get<float>();
126- uint16_t loopRow = calcRowNumSub;126+ uint16_t loopRow = calcRowNumSub;
127- 127+ 
128- constexpr uint32_t oneRepeat = V_LENGTH;128+ constexpr uint32_t oneRepeat = V_LENGTH;
129- int64_t cols = colsAlignBlock_;129+ int64_t cols = colsAlignBlock_;
130- uint16_t repeatCount = DivCeil(cols_, oneRepeat);130+ uint16_t repeatCount = DivCeil(cols_, oneRepeat);
131- __local_mem__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();131+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
132- __local_mem__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();132+ __ubuf__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();
133- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();133+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
134- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();134+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
135- __local_mem__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();135+ __ubuf__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();
136- __VEC_SCOPE__136+ __VEC_SCOPE__
137- {137+ {
138- RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;138+ RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;
139- for (uint16_t r = 0; r < loopRow; r++) {139+ for (uint16_t r = 0; r < loopRow; r++) {
140- uint32_t sreg = cols_;140+ uint32_t sreg = cols_;
141- MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();141+ MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();
142- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));142+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));
143- for (uint16_t i = 0; i < repeatCount; i++) {143+ for (uint16_t i = 0; i < repeatCount; i++) {
144- maskReg = UpdateMask<float>(sreg);144+ maskReg = UpdateMask<float>(sreg);
145- LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);145+ LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
146- LoadAndCast(dyReg, dyAddr, maskReg, r * cols + i * oneRepeat);146+ LoadAndCast(dyReg, dyAddr, maskReg, r * cols + i * oneRepeat);
147- Mul(mulReg2, dyReg, gammaReg, maskReg);147+ Mul(mulReg2, dyReg, gammaReg, maskReg);
148- LoadAndCast(xReg, xAddr, maskReg, r * cols + i * oneRepeat);148+ LoadAndCast(xReg, xAddr, maskReg, r * cols + i * oneRepeat);
149- Mul(mulReg0, xReg, rstdReg, maskReg);149+ Mul(mulReg0, xReg, rstdReg, maskReg);
150- Mul(mulReg3, mulReg2, mulReg0, maskReg);150+ Mul(mulReg3, mulReg2, mulReg0, maskReg);
151- DataCopy(reduceAddr + static_cast<uint32_t>(r * colsAlign2VL_ + i * oneRepeat), mulReg3, maskReg);151+ StoreAlign(reduceAddr + static_cast<uint32_t>(r * colsAlign2VL_ + i * oneRepeat), mulReg3, maskReg);
152- }152+ }
153- }153+ }
154- }154+ }
155- 155+ 
156- MultiReduceSum(tmpSumLocal, reduceLocal, calcRowNumSub);156+ MultiReduceSum(tmpSumLocal, reduceLocal, calcRowNumSub);
157- LocalTensor<T_DX> dxLocal = outQueueDx_.AllocTensor<T_DX>();157+ LocalTensor<T_DX> dxLocal = outQueueDx_.AllocTensor<T_DX>();
158- LocalTensor<T_SCALES_X> scalesXLocal;158+ LocalTensor<T_SCALES_X> scalesXLocal;
159- LocalTensor<T_OFFSET_X> offsetXLocal;159+ LocalTensor<T_OFFSET_X> offsetXLocal;
160- __local_mem__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();160+ __ubuf__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();
161- __local_mem__ T_DX* dxAddr = (__ubuf__ T_DX*)dxLocal.GetPhyAddr();161+ __ubuf__ T_DX* dxAddr = (__ubuf__ T_DX*)dxLocal.GetPhyAddr();
162- __local_mem__ T_SCALES_X* scalesXAddr;162+ __ubuf__ T_SCALES_X* scalesXAddr;
163- __local_mem__ T_OFFSET_X* offsetXAddr;163+ __ubuf__ T_OFFSET_X* offsetXAddr;
164- 164+ 
165- scalesXLocal = scalesXLocal_;165+ scalesXLocal = scalesXLocal_;
166- scalesXAddr = (__ubuf__ T_SCALES_X*)scalesXLocal.GetPhyAddr();166+ scalesXAddr = (__ubuf__ T_SCALES_X*)scalesXLocal.GetPhyAddr();
167- if constexpr (HAS_OFFSET_X) {167+ if constexpr (HAS_OFFSET_X) {
168- offsetXLocal = offsetXLocal_;168+ offsetXLocal = offsetXLocal_;
169- offsetXAddr = (__ubuf__ T_OFFSET_X*)offsetXLocal.GetPhyAddr();169+ offsetXAddr = (__ubuf__ T_OFFSET_X*)offsetXLocal.GetPhyAddr();
170- }170+ }
171- 171+ 
172- __VEC_SCOPE__172+ __VEC_SCOPE__
173- {173+ {
174- RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;174+ RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;
175- RegTensor<float> scalesXReg, scalesXResultReg, offsetXReg;175+ RegTensor<float> scalesXReg, scalesXResultReg, offsetXReg;
176- for (uint16_t r = 0; r < loopRow; r++) {176+ for (uint16_t r = 0; r < loopRow; r++) {
177- uint32_t sreg = cols_;177+ uint32_t sreg = cols_;
178- int64_t cols = colsAlignBlock_;178+ int64_t cols = colsAlignBlock_;
179- int64_t colsAlignHiFP8 = colsAlignHiFP8_;179+ int64_t colsAlignHiFP8 = colsAlignHiFP8_;
180- MaskReg maskReg = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();180+ MaskReg maskReg = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
181- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));181+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + static_cast<uint32_t>(r));
182- DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr + static_cast<uint32_t>(r));182+ LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr + static_cast<uint32_t>(r));
183- Muls(meanReg, meanReg, avgFactor1_, maskReg);183+ Muls(meanReg, meanReg, avgFactor1_, maskReg);
184- for (uint16_t i = 0; i < repeatCount; i++) {184+ for (uint16_t i = 0; i < repeatCount; i++) {
185- maskReg = UpdateMask<float>(sreg);185+ maskReg = UpdateMask<float>(sreg);
186- LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);186+ LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
187- 187+ 
188- LoadAndCast(dyReg, dyAddr, maskReg, r * cols + i * oneRepeat);188+ LoadAndCast(dyReg, dyAddr, maskReg, r * cols + i * oneRepeat);
189- Mul(mulReg2, dyReg, gammaReg, maskReg);189+ Mul(mulReg2, dyReg, gammaReg, maskReg);
190- 190+ 
191- LoadAndCast(xReg, xAddr, maskReg, r * cols + i * oneRepeat);191+ LoadAndCast(xReg, xAddr, maskReg, r * cols + i * oneRepeat);
192- Mul(mulReg0, xReg, rstdReg, maskReg);192+ Mul(mulReg0, xReg, rstdReg, maskReg);
193- Mul(mulReg4, mulReg0, meanReg, maskReg);193+ Mul(mulReg4, mulReg0, meanReg, maskReg);
194- Sub(subReg, mulReg2, mulReg4, maskReg);194+ Sub(subReg, mulReg2, mulReg4, maskReg);
195- Mul(dxReg, subReg, rstdReg, maskReg);195+ Mul(dxReg, subReg, rstdReg, maskReg);
196- // cal quant196+ // cal quant
197- LoadTensorForDtypeTIn(scalesXAddr, scalesXReg, maskReg);197+ LoadTensorForDtypeTIn(scalesXAddr, scalesXReg, maskReg);
198- if constexpr (DIV_MODE) {198+ if constexpr (DIV_MODE) {
199- Div(scalesXResultReg, dxReg, scalesXReg, maskReg);199+ Div(scalesXResultReg, dxReg, scalesXReg, maskReg);
200- } else {200+ } else {
201- Mul(scalesXResultReg, dxReg, scalesXReg, maskReg);201+ Mul(scalesXResultReg, dxReg, scalesXReg, maskReg);
202- }202+ }
203- if constexpr (HAS_OFFSET_X) {203+ if constexpr (HAS_OFFSET_X) {
204- LoadTensorForDtypeTIn(offsetXAddr, offsetXReg, maskReg);204+ LoadTensorForDtypeTIn(offsetXAddr, offsetXReg, maskReg);
205- Add(scalesXResultReg, scalesXResultReg, offsetXReg, maskReg);205+ Add(scalesXResultReg, scalesXResultReg, offsetXReg, maskReg);
206- }206+ }
207- if constexpr (IsSameType<T_DX, hifloat8_t>::value) {207+ if constexpr (IsSameType<T_DX, hifloat8_t>::value) {
208- RegTensor<T_DX> dxRegHif8;208+ RegTensor<T_DX> dxRegHif8;
209- Cast<T_DX, float, castTraitFp322Hifp8>(dxRegHif8, scalesXResultReg, maskReg);209+ Cast<T_DX, float, castTraitFp322Hifp8>(dxRegHif8, scalesXResultReg, maskReg);
210- DataCopy<T_DX, StoreDist::DIST_PACK4_B32>(210+ StoreAlign<T_DX, StoreDist::DIST_PACK4_B32>(
211- dxAddr + static_cast<uint32_t>(r * colsAlignHiFP8 + i * oneRepeat), dxRegHif8, maskReg);211+ dxAddr + static_cast<uint32_t>(r * colsAlignHiFP8 + i * oneRepeat), dxRegHif8, maskReg);
212- } else if constexpr (IsSameType<T_DX, int8_t>::value) {212+ } else if constexpr (IsSameType<T_DX, int8_t>::value) {
213- RegTensor<T_DX> dxRegInt8;213+ RegTensor<T_DX> dxRegInt8;
214- RegTensor<half> dxRegFp16;214+ RegTensor<half> dxRegFp16;
215- RegTensor<int32_t> dxRegInt32;215+ RegTensor<int32_t> dxRegInt32;
216- Cast<int32_t, float, castTraitFp322Int32>(dxRegInt32, scalesXResultReg, maskReg);216+ Cast<int32_t, float, castTraitFp322Int32>(dxRegInt32, scalesXResultReg, maskReg);
217- Cast<float, int32_t, castTraitInt322Fp32>(scalesXResultReg, dxRegInt32, maskReg);217+ Cast<float, int32_t, castTraitInt322Fp32>(scalesXResultReg, dxRegInt32, maskReg);
218- Cast<half, float, castTraitFp322Fp16>(dxRegFp16, scalesXResultReg, maskReg);218+ Cast<half, float, castTraitFp322Fp16>(dxRegFp16, scalesXResultReg, maskReg);
219- Cast<T_DX, half, castTraitFp162Int8>(dxRegInt8, dxRegFp16, maskReg);219+ Cast<T_DX, half, castTraitFp162Int8>(dxRegInt8, dxRegFp16, maskReg);
220- DataCopy<T_DX, StoreDist::DIST_PACK4_B32>(220+ StoreAlign<T_DX, StoreDist::DIST_PACK4_B32>(
221- dxAddr + static_cast<uint32_t>(r * colsAlignHiFP8 + i * oneRepeat), dxRegInt8, maskReg);221+ dxAddr + static_cast<uint32_t>(r * colsAlignHiFP8 + i * oneRepeat), dxRegInt8, maskReg);
222- }222+ }
223- }223+ }
224- }224+ }
225- }225+ }
226- inQueueDy_.FreeTensor(dyLocal);226+ inQueueDy_.FreeTensor(dyLocal);
227- inQueueX_.FreeTensor(xLocal);227+ inQueueX_.FreeTensor(xLocal);
228- inQueueRstd_.FreeTensor(rstdLocal);228+ inQueueRstd_.FreeTensor(rstdLocal);
229- outQueueDx_.EnQue(dxLocal);229+ outQueueDx_.EnQue(dxLocal);
230- CopyOutDx(rowIdx, calcRowNumSub);230+ CopyOutDx(rowIdx, calcRowNumSub);
231- }231+ }
232- 232+ 
233- __aicore__ inline void CopyInRstd(int64_t rowIdx, int64_t count)233+ __aicore__ inline void CopyInRstd(int64_t rowIdx, int64_t count)
234- {234+ {
235- LocalTensor<float> rstdLocal = inQueueRstd_.AllocTensor<float>();235+ LocalTensor<float> rstdLocal = inQueueRstd_.AllocTensor<float>();
236- DataCopyExtParams copyParams{236+ DataCopyExtParams copyParams{
237- 1, // blockCount237+ 1, // blockCount
238- static_cast<uint32_t>(count * sizeof(float)), // blockLen238+ static_cast<uint32_t>(count * sizeof(float)), // blockLen
239- 0, // srcStride239+ 0, // srcStride
240- 0, // dstStride240+ 0, // dstStride
241- 0 // rsv241+ 0 // rsv
242- };242+ };
243- DataCopyPad(rstdLocal, rstdGm_[rowIdx], copyParams, {true, 0, 0, 0});243+ DataCopyPad(rstdLocal, rstdGm_[rowIdx], copyParams, {true, 0, 0, 0});
244- inQueueRstd_.EnQue(rstdLocal);244+ inQueueRstd_.EnQue(rstdLocal);
245- }245+ }
246- 246+ 
247- __aicore__ inline void CopyInGamma()247+ __aicore__ inline void CopyInGamma()
248- {248+ {
249- LocalTensor<T_GAMMA> gammaLocal = inQueueGamma_.AllocTensor<T_GAMMA>();249+ LocalTensor<T_GAMMA> gammaLocal = inQueueGamma_.AllocTensor<T_GAMMA>();
250- DataCopyExtParams copyParams{250+ DataCopyExtParams copyParams{
251- 1, // blockCount251+ 1, // blockCount
252- static_cast<uint32_t>(cols_ * sizeof(T_GAMMA)), // blockLen252+ static_cast<uint32_t>(cols_ * sizeof(T_GAMMA)), // blockLen
253- 0, // srcStride253+ 0, // srcStride
254- 0, // dstStride254+ 0, // dstStride
255- 0 // rsv255+ 0 // rsv
256- };256+ };
257- 257+ 
258- DataCopyPad(gammaLocal, gammaGm_, copyParams, {true, 0, 0, 0});258+ DataCopyPad(gammaLocal, gammaGm_, copyParams, {true, 0, 0, 0});
259- inQueueGamma_.EnQue(gammaLocal);259+ inQueueGamma_.EnQue(gammaLocal);
260- gammaLocal_ = inQueueGamma_.DeQue<T_GAMMA>();260+ gammaLocal_ = inQueueGamma_.DeQue<T_GAMMA>();
261- }261+ }
262- 262+ 
263- __aicore__ inline void CopyInScalesX()263+ __aicore__ inline void CopyInScalesX()
264- {264+ {
265- LocalTensor<T_SCALES_X> scalesXLocal = inQueueScalesX_.AllocTensor<T_SCALES_X>();265+ LocalTensor<T_SCALES_X> scalesXLocal = inQueueScalesX_.AllocTensor<T_SCALES_X>();
266- DataCopyExtParams copyParams{266+ DataCopyExtParams copyParams{
267- 1, // blockCount267+ 1, // blockCount
268- static_cast<uint32_t>(1 * sizeof(T_SCALES_X)), // blockLen268+ static_cast<uint32_t>(1 * sizeof(T_SCALES_X)), // blockLen
269- 0, // srcStride269+ 0, // srcStride
270- 0, // dstStride270+ 0, // dstStride
271- 0 // rsv271+ 0 // rsv
272- };272+ };
273- 273+ 
274- DataCopyPad(scalesXLocal, scalesXGm_, copyParams, {true, 0, 0, 0});274+ DataCopyPad(scalesXLocal, scalesXGm_, copyParams, {true, 0, 0, 0});
275- inQueueScalesX_.EnQue(scalesXLocal);275+ inQueueScalesX_.EnQue(scalesXLocal);
276- scalesXLocal_ = inQueueScalesX_.DeQue<T_SCALES_X>();276+ scalesXLocal_ = inQueueScalesX_.DeQue<T_SCALES_X>();
277- }277+ }
278- 278+ 
279- __aicore__ inline void CopyInOffsetX()279+ __aicore__ inline void CopyInOffsetX()
280- {280+ {
281- LocalTensor<T_OFFSET_X> offsetXLocal = inQueueOffsetX_.AllocTensor<T_OFFSET_X>();281+ LocalTensor<T_OFFSET_X> offsetXLocal = inQueueOffsetX_.AllocTensor<T_OFFSET_X>();
282- DataCopyExtParams copyParams{282+ DataCopyExtParams copyParams{
283- 1, // blockCount283+ 1, // blockCount
284- static_cast<uint32_t>(1 * sizeof(T_OFFSET_X)), // blockLen284+ static_cast<uint32_t>(1 * sizeof(T_OFFSET_X)), // blockLen
285- 0, // srcStride285+ 0, // srcStride
286- 0, // dstStride286+ 0, // dstStride
287- 0 // rsv287+ 0 // rsv
288- };288+ };
289- 289+ 
290- DataCopyPad(offsetXLocal, offsetXGm_, copyParams, {true, 0, 0, 0});290+ DataCopyPad(offsetXLocal, offsetXGm_, copyParams, {true, 0, 0, 0});
291- inQueueOffsetX_.EnQue(offsetXLocal);291+ inQueueOffsetX_.EnQue(offsetXLocal);
292- offsetXLocal_ = inQueueOffsetX_.DeQue<T_OFFSET_X>();292+ offsetXLocal_ = inQueueOffsetX_.DeQue<T_OFFSET_X>();
293- }293+ }
294- 294+ 
295- __aicore__ inline void CopyInDy(int64_t rowIdx, int64_t calcRow)295+ __aicore__ inline void CopyInDy(int64_t rowIdx, int64_t calcRow)
296- {296+ {
297- LocalTensor<T_DY> dyLocal = inQueueDy_.AllocTensor<T_DY>();297+ LocalTensor<T_DY> dyLocal = inQueueDy_.AllocTensor<T_DY>();
298- DataCopyExtParams copyParams{298+ DataCopyExtParams copyParams{
299- static_cast<uint16_t>(calcRow), // blockCount299+ static_cast<uint16_t>(calcRow), // blockCount
300- static_cast<uint32_t>(cols_ * sizeof(T_DY)), // blockLen300+ static_cast<uint32_t>(cols_ * sizeof(T_DY)), // blockLen
301- 0, // srcStride301+ 0, // srcStride
302- 0, // dstStride302+ 0, // dstStride
303- 0 // rsv303+ 0 // rsv
304- };304+ };
305- 305+ 
306- DataCopyPad(dyLocal, dyGm_[rowIdx * cols_], copyParams, {true, 0, 0, 0});306+ DataCopyPad(dyLocal, dyGm_[rowIdx * cols_], copyParams, {true, 0, 0, 0});
307- inQueueDy_.EnQue(dyLocal);307+ inQueueDy_.EnQue(dyLocal);
308- }308+ }
309- 309+ 
310- __aicore__ inline void CopyInX(int64_t rowIdx, int64_t calcRow)310+ __aicore__ inline void CopyInX(int64_t rowIdx, int64_t calcRow)
311- {311+ {
312- LocalTensor<T_X> xLocal = inQueueX_.AllocTensor<T_X>();312+ LocalTensor<T_X> xLocal = inQueueX_.AllocTensor<T_X>();
313- DataCopyExtParams copyParams{313+ DataCopyExtParams copyParams{
314- static_cast<uint16_t>(calcRow), // blockCount314+ static_cast<uint16_t>(calcRow), // blockCount
315- static_cast<uint32_t>(cols_ * sizeof(T_X)), // blockLen315+ static_cast<uint32_t>(cols_ * sizeof(T_X)), // blockLen
316- 0, // srcStride316+ 0, // srcStride
317- 0, // dstStride317+ 0, // dstStride
318- 0 // rsv318+ 0 // rsv
319- };319+ };
320- 320+ 
321- DataCopyPad(xLocal, xGm_[rowIdx * cols_], copyParams, {true, 0, 0, 0});321+ DataCopyPad(xLocal, xGm_[rowIdx * cols_], copyParams, {true, 0, 0, 0});
322- inQueueX_.EnQue(xLocal);322+ inQueueX_.EnQue(xLocal);
323- }323+ }
324- 324+ 
325- __aicore__ inline void MultiReduceSum(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal, int64_t rows)325+ __aicore__ inline void MultiReduceSum(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal, int64_t rows)
326- {326+ {
327- __local_mem__ float* srcAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();327+ __ubuf__ float* srcAddr = (__ubuf__ float*)srcLocal.GetPhyAddr();
328- uint32_t colsTail = colsAlign2VL_ - cols_;328+ uint32_t colsTail = colsAlign2VL_ - cols_;
329- if (colsTail > V_LENGTH) {329+ if (colsTail > V_LENGTH) {
330- // 当要补的个数大于64时，需要两个寄存器进行填充（一个完整的全0 regtensor 加上 利用shiftleft将非对齐位置补0）330+ // 当要补的个数大于64时，需要两个寄存器进行填充（一个完整的全0 regtensor 加上 利用shiftleft将非对齐位置补0）
331- uint32_t colsStartLastTwoVL = colsAlign2VL_ - V_LENGTH * NUM_TWO;331+ uint32_t colsStartLastTwoVL = colsAlign2VL_ - V_LENGTH * NUM_TWO;
332- uint32_t colsStartLastOneVL = colsAlign2VL_ - V_LENGTH;332+ uint32_t colsStartLastOneVL = colsAlign2VL_ - V_LENGTH;
333- uint32_t colsValidLastTwoVL = V_LENGTH * NUM_TWO - colsTail;333+ uint32_t colsValidLastTwoVL = V_LENGTH * NUM_TWO - colsTail;
334- __VEC_SCOPE__334+ __VEC_SCOPE__
335- {335+ {
336- RegTensor<float> xTailReg;336+ RegTensor<float> xTailReg;
337- RegTensor<float> xTailRegshiftLeft;337+ RegTensor<float> xTailRegshiftLeft;
338- RegTensor<float> srcReg;338+ RegTensor<float> srcReg;
339- MaskReg pregTail = UpdateMask<float>(colsValidLastTwoVL);339+ MaskReg pregTail = UpdateMask<float>(colsValidLastTwoVL);
340- MaskReg maskRegAll = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();340+ MaskReg maskRegAll = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
341- Duplicate(srcReg, 0.0f, maskRegAll);341+ Duplicate(srcReg, 0.0f, maskRegAll);
342- for (uint16_t r = 0; r < (uint16_t)rows; r++) {342+ for (uint16_t r = 0; r < (uint16_t)rows; r++) {
343- DataCopy(xTailReg, srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastTwoVL));343+ LoadAlign(xTailReg, srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastTwoVL));
344- // 利用shiftleft将非对齐位置补0344+ // 利用shiftleft将非对齐位置补0
345- ShiftLefts((RegTensor<uint32_t>&)xTailRegshiftLeft, (RegTensor<uint32_t>&)xTailReg,345+ ShiftLefts((RegTensor<uint32_t>&)xTailRegshiftLeft, (RegTensor<uint32_t>&)xTailReg,
346- static_cast<int16_t>(0), pregTail);346+ static_cast<int16_t>(0), pregTail);
347- DataCopy(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastTwoVL), xTailRegshiftLeft,347+ StoreAlign(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastTwoVL),
348- maskRegAll);348+ xTailRegshiftLeft, maskRegAll);
349- DataCopy(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL), srcReg,349+ StoreAlign(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL), srcReg,
350- maskRegAll);350+ maskRegAll);
351- }351+ }
352- }352+ }
353- } else if (colsTail == V_LENGTH) {353+ } else if (colsTail == V_LENGTH) {
354- // 当要补的个数等于64时，直接dup一个全0的regtensor进行填充354+ // 当要补的个数等于64时，直接dup一个全0的regtensor进行填充
355- uint32_t colsStartLastOneVL = colsAlign2VL_ - V_LENGTH;355+ uint32_t colsStartLastOneVL = colsAlign2VL_ - V_LENGTH;
356- __VEC_SCOPE__356+ __VEC_SCOPE__
357- {357+ {
358- RegTensor<float> srcReg;358+ RegTensor<float> srcReg;
359- MaskReg maskRegAll = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();359+ MaskReg maskRegAll = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
360- Duplicate(srcReg, 0.0f, maskRegAll);360+ Duplicate(srcReg, 0.0f, maskRegAll);
361- for (uint16_t r = 0; r < (uint16_t)rows; r++) {361+ for (uint16_t r = 0; r < (uint16_t)rows; r++) {
362- DataCopy(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL), srcReg,362+ StoreAlign(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL), srcReg,
363- maskRegAll);363+ maskRegAll);
364- }364+ }
365- }365+ }
366- } else if (colsTail > 0) {366+ } else if (colsTail > 0) {
367- // 当要补的个数小于64时，利用shiftleft将非对齐位置补0367+ // 当要补的个数小于64时，利用shiftleft将非对齐位置补0
368- uint32_t colsStartLastOneVL = colsAlign2VL_ - V_LENGTH;368+ uint32_t colsStartLastOneVL = colsAlign2VL_ - V_LENGTH;
369- uint32_t colsValidLastOneVL = V_LENGTH - colsTail;369+ uint32_t colsValidLastOneVL = V_LENGTH - colsTail;
370- __VEC_SCOPE__370+ __VEC_SCOPE__
371- {371+ {
372- RegTensor<float> xTailReg;372+ RegTensor<float> xTailReg;
373- RegTensor<float> xTailRegshiftLeft;373+ RegTensor<float> xTailRegshiftLeft;
374- MaskReg pregTail = UpdateMask<float>(colsValidLastOneVL);374+ MaskReg pregTail = UpdateMask<float>(colsValidLastOneVL);
375- MaskReg maskRegAll = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();375+ MaskReg maskRegAll = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
376- for (uint16_t r = 0; r < (uint16_t)rows; r++) {376+ for (uint16_t r = 0; r < (uint16_t)rows; r++) {
377- DataCopy(xTailReg, srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL));377+ LoadAlign(xTailReg, srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL));
378- // 利用shiftleft将非对齐位置补0378+ // 利用shiftleft将非对齐位置补0
379- ShiftLefts((RegTensor<uint32_t>&)xTailRegshiftLeft, (RegTensor<uint32_t>&)xTailReg,379+ ShiftLefts((RegTensor<uint32_t>&)xTailRegshiftLeft, (RegTensor<uint32_t>&)xTailReg,
380- static_cast<int16_t>(0), pregTail);380+ static_cast<int16_t>(0), pregTail);
381- DataCopy(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL), xTailRegshiftLeft,381+ StoreAlign(srcAddr + static_cast<uint32_t>(r * colsAlign2VL_ + colsStartLastOneVL),
382- maskRegAll);382+ xTailRegshiftLeft, maskRegAll);
383- }383+ }
384- }384+ }
385- }385+ }
386- uint32_t srcShape[2] = {uint32_t(rows), uint32_t(colsAlign2VL_)};386+ uint32_t srcShape[2] = {uint32_t(rows), uint32_t(colsAlign2VL_)};
387- AscendC::ReduceSum<float, AscendC::Pattern::Reduce::AR, true>(dstLocal, srcLocal, srcShape, false);387+ AscendC::ReduceSum<float, AscendC::Pattern::Reduce::AR, true>(dstLocal, srcLocal, srcShape, false);
388- }388+ }
389- 389+ 
390- template <typename T_IN>390+ template <typename T_IN>
391- __aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg)391+ __aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_IN* src, RegTensor<float>& dst, MaskReg& preg)
392- {392+ {
393- if constexpr (IsSameType<T_IN, float>::value) {393+ if constexpr (IsSameType<T_IN, float>::value) {
394- DataCopy<float, LoadDist::DIST_BRC_B32>(dst, src);394+ LoadAlign<float, LoadDist::DIST_BRC_B32>(dst, src);
395- } else if constexpr (IsSameType<T_IN, int32_t>::value) {395+ } else if constexpr (IsSameType<T_IN, int32_t>::value) {
396- RegTensor<T_IN> xIn;396+ RegTensor<T_IN> xIn;
397- DataCopy<int32_t, LoadDist::DIST_BRC_B32>(xIn, src);397+ LoadAlign<int32_t, LoadDist::DIST_BRC_B32>(xIn, src);
398- Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);398+ Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);
399- } else {399+ } else {
400- RegTensor<T_IN> xIn;400+ RegTensor<T_IN> xIn;
401- DataCopy<T_IN, LoadDist::DIST_BRC_B16>(xIn, src);401+ LoadAlign<T_IN, LoadDist::DIST_BRC_B16>(xIn, src);
402- Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg);402+ Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg);
403- }403+ }
404- }404+ }
405- 405+ 
406- __aicore__ inline void CopyOutDx(int64_t rowIdx, int64_t calcRow)406+ __aicore__ inline void CopyOutDx(int64_t rowIdx, int64_t calcRow)
407- {407+ {
408- LocalTensor<T_DX> dxLocal = outQueueDx_.DeQue<T_DX>();408+ LocalTensor<T_DX> dxLocal = outQueueDx_.DeQue<T_DX>();
409- DataCopyExtParams copyParams{409+ DataCopyExtParams copyParams{
410- static_cast<uint16_t>(calcRow), // blockCount410+ static_cast<uint16_t>(calcRow), // blockCount
411- static_cast<uint32_t>(cols_ * sizeof(T_DX)), // blockLen411+ static_cast<uint32_t>(cols_ * sizeof(T_DX)), // blockLen
412- 0, // srcStride412+ 0, // srcStride
413- 0, // dstStride413+ 0, // dstStride
414- 0 // rsv414+ 0 // rsv
415- };415+ };
416- DataCopyPad(dxGm_[rowIdx * cols_], dxLocal, copyParams);416+ DataCopyPad(dxGm_[rowIdx * cols_], dxLocal, copyParams);
417- outQueueDx_.FreeTensor(dxLocal);417+ outQueueDx_.FreeTensor(dxLocal);
418- }418+ }
419- 419+ 
420-private:420+private:
421- TPipe* Ppipe_;421+ TPipe* Ppipe_;
422- const RmsNormGradQuantRegbaseDxTilingData* tiling_;422+ const RmsNormGradQuantRegbaseDxTilingData* tiling_;
423- GlobalTensor<T_DY> dyGm_;423+ GlobalTensor<T_DY> dyGm_;
424- GlobalTensor<T_X> xGm_;424+ GlobalTensor<T_X> xGm_;
425- GlobalTensor<T_GAMMA> gammaGm_;425+ GlobalTensor<T_GAMMA> gammaGm_;
426- GlobalTensor<float> rstdGm_;426+ GlobalTensor<float> rstdGm_;
427- GlobalTensor<T_DX> dxGm_;427+ GlobalTensor<T_DX> dxGm_;
428- GlobalTensor<T_SCALES_X> scalesXGm_;428+ GlobalTensor<T_SCALES_X> scalesXGm_;
429- GlobalTensor<T_OFFSET_X> offsetXGm_;429+ GlobalTensor<T_OFFSET_X> offsetXGm_;
430- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueDy_;430+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueDy_;
431- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueX_;431+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueX_;
432- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueRstd_;432+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueRstd_;
433- TQue<QuePosition::VECOUT, DEPTH_TWO> outQueueDx_;433+ TQue<QuePosition::VECOUT, DEPTH_TWO> outQueueDx_;
434- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueGamma_;434+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueGamma_;
435- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueScalesX_;435+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueScalesX_;
436- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueOffsetX_;436+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueOffsetX_;
437- TBuf<TPosition::VECCALC> reduceBuf_;437+ TBuf<TPosition::VECCALC> reduceBuf_;
438- TBuf<TPosition::VECCALC> tmpSumBuf_;438+ TBuf<TPosition::VECCALC> tmpSumBuf_;
439- LocalTensor<T_GAMMA> gammaLocal_;439+ LocalTensor<T_GAMMA> gammaLocal_;
440- LocalTensor<T_SCALES_X> scalesXLocal_;440+ LocalTensor<T_SCALES_X> scalesXLocal_;
441- LocalTensor<T_OFFSET_X> offsetXLocal_;441+ LocalTensor<T_OFFSET_X> offsetXLocal_;
442- 442+ 
443- uint32_t usedCoreNum_;443+ uint32_t usedCoreNum_;
444- int64_t rows_;444+ int64_t rows_;
445- int64_t cols_;445+ int64_t cols_;
446- int64_t colsAlignBlock_;446+ int64_t colsAlignBlock_;
447- int64_t colsAlign2VL_;447+ int64_t colsAlign2VL_;
448- int64_t colsAlignHiFP8_;448+ int64_t colsAlignHiFP8_;
449- int64_t blockFactor_;449+ int64_t blockFactor_;
450- int64_t ubFactor_;450+ int64_t ubFactor_;
451- int64_t ubFactorN_;451+ int64_t ubFactorN_;
452- int64_t ubFactorD_;452+ int64_t ubFactorD_;
453- float avgFactor1_;453+ float avgFactor1_;
454-};454+};
455-} // namespace RmsNormGradQuant455+} // namespace RmsNormGradQuant
456-#endif // RMS_NORM_GRAD_REGBASE_DX_FULL_LOAD_H456+#endif // RMS_NORM_GRAD_REGBASE_DX_FULL_LOAD_H
Mnorm/rms_norm_grad_quant/op_kernel/arch35/rms_norm_grad_quant_dx_split_d.h+463-463
@@ -1,463 +1,463 @@
1-/**1+/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/*!11+/*!
12- * \file rms_norm_grad_regbase_dx_split_d.h12+ * \file rms_norm_grad_regbase_dx_split_d.h
13- * \brief RmsNormGrad Regbase DX Split D kernel File13+ * \brief RmsNormGrad Regbase DX Split D kernel File
14- */14+ */
15- 15+ 
16-#ifndef RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H16+#ifndef RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H
17-#define RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H17+#define RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H
18- 18+ 
19-#include "kernel_tiling/kernel_tiling.h"19+#include "kernel_tiling/kernel_tiling.h"
20-#include "kernel_operator.h"20+#include "kernel_operator.h"
21-#include "rms_norm_grad_quant_common.h"21+#include "rms_norm_grad_quant_common.h"
22- 22+ 
23-namespace RmsNormGradQuant {23+namespace RmsNormGradQuant {
24-using namespace AscendC;24+using namespace AscendC;
25-template <typename T_DY, typename T_X, typename T_GAMMA, typename T_DX, typename T_DGAMMA, typename T_SCALES_X,25+template <typename T_DY, typename T_X, typename T_GAMMA, typename T_DX, typename T_DGAMMA, typename T_SCALES_X,
26- typename T_OFFSET_X, bool HAS_OFFSET_X, bool DIV_MODE>26+ typename T_OFFSET_X, bool HAS_OFFSET_X, bool DIV_MODE>
27-class RegbaseDxSplitD {27+class RegbaseDxSplitD {
28-public:28+public:
29- __aicore__ inline RegbaseDxSplitD(TPipe* pipe, const RmsNormGradQuantRegbaseDxTilingData* tilingData)29+ __aicore__ inline RegbaseDxSplitD(TPipe* pipe, const RmsNormGradQuantRegbaseDxTilingData* tilingData)
30- : Ppipe_(pipe), tiling_(tilingData)30+ : Ppipe_(pipe), tiling_(tilingData)
31- {}31+ {}
32- 32+ 
33- __aicore__ inline void Init(__gm__ uint8_t* dy, __gm__ uint8_t* x, __gm__ uint8_t* rstd, __gm__ uint8_t* gamma,33+ __aicore__ inline void Init(__gm__ uint8_t* dy, __gm__ uint8_t* x, __gm__ uint8_t* rstd, __gm__ uint8_t* gamma,
34- __gm__ uint8_t* scales_x, __gm__ uint8_t* offset_x, __gm__ uint8_t* dx,34+ __gm__ uint8_t* scales_x, __gm__ uint8_t* offset_x, __gm__ uint8_t* dx,
35- __gm__ uint8_t* dgamma)35+ __gm__ uint8_t* dgamma)
36- {36+ {
37-#if (__NPU_ARCH__ == 3510)37+#if (__NPU_ARCH__ == 3510)
38- if constexpr (IsSameType<T_DX, hifloat8_t>::value) {38+ if constexpr (IsSameType<T_DX, hifloat8_t>::value) {
39- AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);39+ AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
40- }40+ }
41-#endif41+#endif
42- usedCoreNum_ = tiling_->usedCoreNumDx;42+ usedCoreNum_ = tiling_->usedCoreNumDx;
43- uint32_t coreIdx = GetBlockIdx();43+ uint32_t coreIdx = GetBlockIdx();
44- if (coreIdx >= usedCoreNum_) {44+ if (coreIdx >= usedCoreNum_) {
45- return;45+ return;
46- }46+ }
47- rows_ = tiling_->rows;47+ rows_ = tiling_->rows;
48- cols_ = tiling_->cols;48+ cols_ = tiling_->cols;
49- blockFactor_ = tiling_->blockFactorDx; // ceilDiv(rows_, usedCoreNum)49+ blockFactor_ = tiling_->blockFactorDx; // ceilDiv(rows_, usedCoreNum)
50- ubFactorD_ = UB_FACTOR_DX_SPLIT_D; // 固定值50+ ubFactorD_ = UB_FACTOR_DX_SPLIT_D; // 固定值
51- bodyPart_ = tiling_->bodyPart; // 小于cols的最大二次幂51+ bodyPart_ = tiling_->bodyPart; // 小于cols的最大二次幂
52- avgFactor1_ = 1.0f / cols_;52+ avgFactor1_ = 1.0f / cols_;
53- dyGm_.SetGlobalBuffer((__gm__ T_DY*)dy + coreIdx * blockFactor_ * cols_);53+ dyGm_.SetGlobalBuffer((__gm__ T_DY*)dy + coreIdx * blockFactor_ * cols_);
54- xGm_.SetGlobalBuffer((__gm__ T_X*)x + coreIdx * blockFactor_ * cols_);54+ xGm_.SetGlobalBuffer((__gm__ T_X*)x + coreIdx * blockFactor_ * cols_);
55- rstdGm_.SetGlobalBuffer((__gm__ float*)rstd + coreIdx * blockFactor_);55+ rstdGm_.SetGlobalBuffer((__gm__ float*)rstd + coreIdx * blockFactor_);
56- gammaGm_.SetGlobalBuffer((__gm__ T_GAMMA*)gamma);56+ gammaGm_.SetGlobalBuffer((__gm__ T_GAMMA*)gamma);
57- dxGm_.SetGlobalBuffer((__gm__ T_DX*)dx + coreIdx * blockFactor_ * cols_);57+ dxGm_.SetGlobalBuffer((__gm__ T_DX*)dx + coreIdx * blockFactor_ * cols_);
58- 58+ 
59- Ppipe_->InitBuffer(inQueueDy_, DB_NUM, ubFactorD_ * sizeof(float));59+ Ppipe_->InitBuffer(inQueueDy_, DB_NUM, ubFactorD_ * sizeof(float));
60- Ppipe_->InitBuffer(inQueueX_, DB_NUM, ubFactorD_ * sizeof(float));60+ Ppipe_->InitBuffer(inQueueX_, DB_NUM, ubFactorD_ * sizeof(float));
61- Ppipe_->InitBuffer(inQueueRstd_, DB_NUM, V_LENGTH * sizeof(float));61+ Ppipe_->InitBuffer(inQueueRstd_, DB_NUM, V_LENGTH * sizeof(float));
62- Ppipe_->InitBuffer(inQueueGamma_, DB_NUM, ubFactorD_ * sizeof(float));62+ Ppipe_->InitBuffer(inQueueGamma_, DB_NUM, ubFactorD_ * sizeof(float));
63- Ppipe_->InitBuffer(outQueueDx_, DB_NUM, ubFactorD_ * sizeof(float));63+ Ppipe_->InitBuffer(outQueueDx_, DB_NUM, ubFactorD_ * sizeof(float));
64- Ppipe_->InitBuffer(reduceBuf_, DB_NUM * ubFactorD_ * sizeof(float));64+ Ppipe_->InitBuffer(reduceBuf_, DB_NUM * ubFactorD_ * sizeof(float));
65- Ppipe_->InitBuffer(level0Buf_, ONCE_VECTOR_SIZE * sizeof(float));65+ Ppipe_->InitBuffer(level0Buf_, ONCE_VECTOR_SIZE * sizeof(float));
66- Ppipe_->InitBuffer(level1Buf_, ONCE_VECTOR_SIZE * sizeof(float));66+ Ppipe_->InitBuffer(level1Buf_, ONCE_VECTOR_SIZE * sizeof(float));
67- Ppipe_->InitBuffer(level2Buf_, ONCE_VECTOR_SIZE * sizeof(float));67+ Ppipe_->InitBuffer(level2Buf_, ONCE_VECTOR_SIZE * sizeof(float));
68- Ppipe_->InitBuffer(tmpSumBuf_, V_LENGTH * sizeof(float));68+ Ppipe_->InitBuffer(tmpSumBuf_, V_LENGTH * sizeof(float));
69- Ppipe_->InitBuffer(workBuf_, ONCE_VECTOR_SIZE * sizeof(float));69+ Ppipe_->InitBuffer(workBuf_, ONCE_VECTOR_SIZE * sizeof(float));
70- scalesXGm_.SetGlobalBuffer((__gm__ T_SCALES_X*)scales_x);70+ scalesXGm_.SetGlobalBuffer((__gm__ T_SCALES_X*)scales_x);
71- Ppipe_->InitBuffer(scalesXBuf_, sizeof(T_SCALES_X));71+ Ppipe_->InitBuffer(scalesXBuf_, sizeof(T_SCALES_X));
72- if constexpr (HAS_OFFSET_X) {72+ if constexpr (HAS_OFFSET_X) {
73- offsetXGm_.SetGlobalBuffer((__gm__ T_OFFSET_X*)offset_x);73+ offsetXGm_.SetGlobalBuffer((__gm__ T_OFFSET_X*)offset_x);
74- Ppipe_->InitBuffer(offsetXBuf_, sizeof(T_OFFSET_X));74+ Ppipe_->InitBuffer(offsetXBuf_, sizeof(T_OFFSET_X));
75- }75+ }
76- }76+ }
77- __aicore__ inline void Process()77+ __aicore__ inline void Process()
78- {78+ {
79- uint32_t coreIdx = GetBlockIdx();79+ uint32_t coreIdx = GetBlockIdx();
80- if (coreIdx >= usedCoreNum_) {80+ if (coreIdx >= usedCoreNum_) {
81- return;81+ return;
82- }82+ }
83- // copyInScalesX83+ // copyInScalesX
84- CopyInScalesX();84+ CopyInScalesX();
85- if constexpr (HAS_OFFSET_X) {85+ if constexpr (HAS_OFFSET_X) {
86- CopyInOffsetX();86+ CopyInOffsetX();
87- }87+ }
88- int64_t blockTail = rows_ - (usedCoreNum_ - 1) * blockFactor_;88+ int64_t blockTail = rows_ - (usedCoreNum_ - 1) * blockFactor_;
89- int64_t calcRowNum = coreIdx == usedCoreNum_ - 1 ? blockTail : blockFactor_;89+ int64_t calcRowNum = coreIdx == usedCoreNum_ - 1 ? blockTail : blockFactor_;
90- for (int64_t rowIdx = 0; rowIdx < calcRowNum; rowIdx++) {90+ for (int64_t rowIdx = 0; rowIdx < calcRowNum; rowIdx++) {
91- SubProcess(rowIdx);91+ SubProcess(rowIdx);
92- }92+ }
93- }93+ }
94- 94+ 
95- __aicore__ inline void SubProcess(int64_t rowIdx)95+ __aicore__ inline void SubProcess(int64_t rowIdx)
96- {96+ {
97- CopyInRstd(rowIdx, 1);97+ CopyInRstd(rowIdx, 1);
98- LocalTensor<float> rstdLocal = inQueueRstd_.DeQue<float>();98+ LocalTensor<float> rstdLocal = inQueueRstd_.DeQue<float>();
99- FormerProcess(rstdLocal, rowIdx);99+ FormerProcess(rstdLocal, rowIdx);
100- LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();100+ LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();
101- Muls(tmpSumLocal, tmpSumLocal, avgFactor1_, 1);101+ Muls(tmpSumLocal, tmpSumLocal, avgFactor1_, 1);
102- LatterProcess(rstdLocal, rowIdx);102+ LatterProcess(rstdLocal, rowIdx);
103- inQueueRstd_.FreeTensor(rstdLocal);103+ inQueueRstd_.FreeTensor(rstdLocal);
104- }104+ }
105- 105+ 
106- __aicore__ inline void FormerProcess(LocalTensor<float>& rstdLocal, int64_t rowIdx)106+ __aicore__ inline void FormerProcess(LocalTensor<float>& rstdLocal, int64_t rowIdx)
107- {107+ {
108- uint32_t level0Offset = 0;108+ uint32_t level0Offset = 0;
109- uint32_t level1Offset = 0;109+ uint32_t level1Offset = 0;
110- uint32_t level2Offset = 0;110+ uint32_t level2Offset = 0;
111- InitLevelLocal();111+ InitLevelLocal();
112- for (int64_t colIdx = 0; colIdx < bodyPart_; colIdx += ubFactorD_) {112+ for (int64_t colIdx = 0; colIdx < bodyPart_; colIdx += ubFactorD_) {
113- CopyInGamma(colIdx, ubFactorD_);113+ CopyInGamma(colIdx, ubFactorD_);
114- CopyInDy(rowIdx, colIdx, ubFactorD_);114+ CopyInDy(rowIdx, colIdx, ubFactorD_);
115- CopyInX(rowIdx, colIdx, ubFactorD_);115+ CopyInX(rowIdx, colIdx, ubFactorD_);
116- ComputeMul<true>(rstdLocal, ubFactorD_);116+ ComputeMul<true>(rstdLocal, ubFactorD_);
117- int64_t tailCount = 0;117+ int64_t tailCount = 0;
118- if (bodyPart_ + colIdx < cols_) {118+ if (bodyPart_ + colIdx < cols_) {
119- int64_t remainCount = cols_ - bodyPart_ - colIdx; // must > 0119+ int64_t remainCount = cols_ - bodyPart_ - colIdx; // must > 0
120- tailCount = Min(remainCount, ubFactorD_);120+ tailCount = Min(remainCount, ubFactorD_);
121- CopyInGamma(bodyPart_ + colIdx, tailCount);121+ CopyInGamma(bodyPart_ + colIdx, tailCount);
122- CopyInDy(rowIdx, bodyPart_ + colIdx, tailCount);122+ CopyInDy(rowIdx, bodyPart_ + colIdx, tailCount);
123- CopyInX(rowIdx, bodyPart_ + colIdx, tailCount);123+ CopyInX(rowIdx, bodyPart_ + colIdx, tailCount);
124- ComputeMul<false>(rstdLocal, tailCount);124+ ComputeMul<false>(rstdLocal, tailCount);
125- }125+ }
126- int64_t reduceCount = ubFactorD_ + tailCount; // [ubFactorD, 2*ubFactorD_]126+ int64_t reduceCount = ubFactorD_ + tailCount; // [ubFactorD, 2*ubFactorD_]
127- ComputeIntoMultiLevel(reduceCount, level0Offset, level1Offset, level2Offset);127+ ComputeIntoMultiLevel(reduceCount, level0Offset, level1Offset, level2Offset);
128- }128+ }
129- FinalLevelReduce(level0Offset, level1Offset);129+ FinalLevelReduce(level0Offset, level1Offset);
130- }130+ }
131- 131+ 
132- __aicore__ inline void InitLevelLocal()132+ __aicore__ inline void InitLevelLocal()
133- {133+ {
134- LocalTensor<float> level0Local = level0Buf_.Get<float>();134+ LocalTensor<float> level0Local = level0Buf_.Get<float>();
135- LocalTensor<float> level1Local = level1Buf_.Get<float>();135+ LocalTensor<float> level1Local = level1Buf_.Get<float>();
136- LocalTensor<float> level2Local = level2Buf_.Get<float>();136+ LocalTensor<float> level2Local = level2Buf_.Get<float>();
137- 137+ 
138- Duplicate(level0Local, 0.0f, ONCE_VECTOR_SIZE);138+ Duplicate(level0Local, 0.0f, ONCE_VECTOR_SIZE);
139- Duplicate(level1Local, 0.0f, ONCE_VECTOR_SIZE);139+ Duplicate(level1Local, 0.0f, ONCE_VECTOR_SIZE);
140- Duplicate(level2Local, 0.0f, ONCE_VECTOR_SIZE);140+ Duplicate(level2Local, 0.0f, ONCE_VECTOR_SIZE);
141- }141+ }
142- 142+ 
143- __aicore__ inline void ComputeIntoMultiLevel(int64_t count, uint32_t& level0Offset, uint32_t& level1Offset,143+ __aicore__ inline void ComputeIntoMultiLevel(int64_t count, uint32_t& level0Offset, uint32_t& level1Offset,
144- uint32_t& level2Offset)144+ uint32_t& level2Offset)
145- {145+ {
146- LocalTensor<float> level0Local = level0Buf_.Get<float>();146+ LocalTensor<float> level0Local = level0Buf_.Get<float>();
147- LocalTensor<float> level1Local = level1Buf_.Get<float>();147+ LocalTensor<float> level1Local = level1Buf_.Get<float>();
148- LocalTensor<float> level2Local = level2Buf_.Get<float>();148+ LocalTensor<float> level2Local = level2Buf_.Get<float>();
149- LocalTensor<float> reduceLocal = reduceBuf_.Get<float>();149+ LocalTensor<float> reduceLocal = reduceBuf_.Get<float>();
150- WholeReduceSum(level0Local, reduceLocal, count, level0Offset);150+ WholeReduceSum(level0Local, reduceLocal, count, level0Offset);
151- level0Offset++;151+ level0Offset++;
152- ComputeMultiLevelReduce(level0Local, level1Local, level2Local, level0Offset, level1Offset, level2Offset);152+ ComputeMultiLevelReduce(level0Local, level1Local, level2Local, level0Offset, level1Offset, level2Offset);
153- }153+ }
154- 154+ 
155- __aicore__ inline void FinalLevelReduce(uint32_t& level0Offset, uint32_t& level1Offset)155+ __aicore__ inline void FinalLevelReduce(uint32_t& level0Offset, uint32_t& level1Offset)
156- {156+ {
157- LocalTensor<float> level0Local = level0Buf_.Get<float>();157+ LocalTensor<float> level0Local = level0Buf_.Get<float>();
158- LocalTensor<float> level1Local = level1Buf_.Get<float>();158+ LocalTensor<float> level1Local = level1Buf_.Get<float>();
159- LocalTensor<float> level2Local = level2Buf_.Get<float>();159+ LocalTensor<float> level2Local = level2Buf_.Get<float>();
160- LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();160+ LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();
161- ComputeMultiLevelMean(tmpSumLocal, 0, level0Local, level1Local, level2Local, level0Offset, level1Offset);161+ ComputeMultiLevelMean(tmpSumLocal, 0, level0Local, level1Local, level2Local, level0Offset, level1Offset);
162- }162+ }
163- 163+ 
164- __aicore__ inline void CopyInRstd(int64_t rowIdx, int64_t count)164+ __aicore__ inline void CopyInRstd(int64_t rowIdx, int64_t count)
165- {165+ {
166- LocalTensor<float> rstdLocal = inQueueRstd_.AllocTensor<float>();166+ LocalTensor<float> rstdLocal = inQueueRstd_.AllocTensor<float>();
167- DataCopyExtParams copyParams{167+ DataCopyExtParams copyParams{
168- 1, // blockCount168+ 1, // blockCount
169- static_cast<uint32_t>(count * sizeof(float)), // blockLen169+ static_cast<uint32_t>(count * sizeof(float)), // blockLen
170- 0, // srcStride170+ 0, // srcStride
171- 0, // dstStride171+ 0, // dstStride
172- 0 // rsv172+ 0 // rsv
173- };173+ };
174- DataCopyPad(rstdLocal, rstdGm_[rowIdx], copyParams, {true, 0, 0, 0});174+ DataCopyPad(rstdLocal, rstdGm_[rowIdx], copyParams, {true, 0, 0, 0});
175- inQueueRstd_.EnQue(rstdLocal);175+ inQueueRstd_.EnQue(rstdLocal);
176- }176+ }
177- 177+ 
178- __aicore__ inline void CopyInGamma(int64_t colIdx, int64_t count)178+ __aicore__ inline void CopyInGamma(int64_t colIdx, int64_t count)
179- {179+ {
180- LocalTensor<T_GAMMA> gammaLocal = inQueueGamma_.AllocTensor<T_GAMMA>();180+ LocalTensor<T_GAMMA> gammaLocal = inQueueGamma_.AllocTensor<T_GAMMA>();
181- DataCopyExtParams copyParams{181+ DataCopyExtParams copyParams{
182- 1, // blockCount182+ 1, // blockCount
183- static_cast<uint32_t>(count * sizeof(T_GAMMA)), // blockLen183+ static_cast<uint32_t>(count * sizeof(T_GAMMA)), // blockLen
184- 0, // srcStride184+ 0, // srcStride
185- 0, // dstStride185+ 0, // dstStride
186- 0 // rsv186+ 0 // rsv
187- };187+ };
188- DataCopyPad(gammaLocal, gammaGm_[colIdx], copyParams, {true, 0, 0, 0});188+ DataCopyPad(gammaLocal, gammaGm_[colIdx], copyParams, {true, 0, 0, 0});
189- inQueueGamma_.EnQue(gammaLocal);189+ inQueueGamma_.EnQue(gammaLocal);
190- }190+ }
191- 191+ 
192- __aicore__ inline void CopyInDy(int64_t rowIdx, int64_t colIdx, int64_t count)192+ __aicore__ inline void CopyInDy(int64_t rowIdx, int64_t colIdx, int64_t count)
193- {193+ {
194- LocalTensor<T_DY> dyLocal = inQueueDy_.AllocTensor<T_DY>();194+ LocalTensor<T_DY> dyLocal = inQueueDy_.AllocTensor<T_DY>();
195- DataCopyExtParams copyParams{195+ DataCopyExtParams copyParams{
196- 1, // blockCount196+ 1, // blockCount
197- static_cast<uint32_t>(count * sizeof(T_DY)), // blockLen197+ static_cast<uint32_t>(count * sizeof(T_DY)), // blockLen
198- 0, // srcStride198+ 0, // srcStride
199- 0, // dstStride199+ 0, // dstStride
200- 0 // rsv200+ 0 // rsv
201- };201+ };
202- DataCopyPad(dyLocal, dyGm_[rowIdx * cols_ + colIdx], copyParams, {true, 0, 0, 0});202+ DataCopyPad(dyLocal, dyGm_[rowIdx * cols_ + colIdx], copyParams, {true, 0, 0, 0});
203- inQueueDy_.EnQue(dyLocal);203+ inQueueDy_.EnQue(dyLocal);
204- }204+ }
205- 205+ 
206- __aicore__ inline void CopyInX(int64_t rowIdx, int64_t colIdx, int64_t count)206+ __aicore__ inline void CopyInX(int64_t rowIdx, int64_t colIdx, int64_t count)
207- {207+ {
208- LocalTensor<T_X> xLocal = inQueueX_.AllocTensor<T_X>();208+ LocalTensor<T_X> xLocal = inQueueX_.AllocTensor<T_X>();
209- DataCopyExtParams copyParams{209+ DataCopyExtParams copyParams{
210- 1, // blockCount210+ 1, // blockCount
211- static_cast<uint32_t>(count * sizeof(T_X)), // blockLen211+ static_cast<uint32_t>(count * sizeof(T_X)), // blockLen
212- 0, // srcStride212+ 0, // srcStride
213- 0, // dstStride213+ 0, // dstStride
214- 0 // rsv214+ 0 // rsv
215- };215+ };
216- DataCopyPad(xLocal, xGm_[rowIdx * cols_ + colIdx], copyParams, {true, 0, 0, 0});216+ DataCopyPad(xLocal, xGm_[rowIdx * cols_ + colIdx], copyParams, {true, 0, 0, 0});
217- inQueueX_.EnQue(xLocal);217+ inQueueX_.EnQue(xLocal);
218- }218+ }
219- 219+ 
220- __aicore__ inline void CopyInScalesX()220+ __aicore__ inline void CopyInScalesX()
221- {221+ {
222- LocalTensor<T_SCALES_X> scalesXLocal = scalesXBuf_.Get<T_SCALES_X>();222+ LocalTensor<T_SCALES_X> scalesXLocal = scalesXBuf_.Get<T_SCALES_X>();
223- DataCopyExtParams copyParams{223+ DataCopyExtParams copyParams{
224- 1, // blockCount224+ 1, // blockCount
225- static_cast<uint32_t>(1 * sizeof(T_SCALES_X)), // blockLen225+ static_cast<uint32_t>(1 * sizeof(T_SCALES_X)), // blockLen
226- 0, // srcStride226+ 0, // srcStride
227- 0, // dstStride227+ 0, // dstStride
228- 0 // rsv228+ 0 // rsv
229- };229+ };
230- 230+ 
231- DataCopyPad(scalesXLocal, scalesXGm_, copyParams, {true, 0, 0, 0});231+ DataCopyPad(scalesXLocal, scalesXGm_, copyParams, {true, 0, 0, 0});
232- }232+ }
233- 233+ 
234- __aicore__ inline void CopyInOffsetX()234+ __aicore__ inline void CopyInOffsetX()
235- {235+ {
236- LocalTensor<T_OFFSET_X> offsetXLocal = offsetXBuf_.Get<T_OFFSET_X>();236+ LocalTensor<T_OFFSET_X> offsetXLocal = offsetXBuf_.Get<T_OFFSET_X>();
237- DataCopyExtParams copyParams{237+ DataCopyExtParams copyParams{
238- 1, // blockCount238+ 1, // blockCount
239- static_cast<uint32_t>(1 * sizeof(T_OFFSET_X)), // blockLen239+ static_cast<uint32_t>(1 * sizeof(T_OFFSET_X)), // blockLen
240- 0, // srcStride240+ 0, // srcStride
241- 0, // dstStride241+ 0, // dstStride
242- 0 // rsv242+ 0 // rsv
243- };243+ };
244- 244+ 
245- DataCopyPad(offsetXLocal, offsetXGm_, copyParams, {true, 0, 0, 0});245+ DataCopyPad(offsetXLocal, offsetXGm_, copyParams, {true, 0, 0, 0});
246- }246+ }
247- 247+ 
248- template <typename T_IN>248+ template <typename T_IN>
249- __aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg)249+ __aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_IN* src, RegTensor<float>& dst, MaskReg& preg)
250- {250+ {
251- if constexpr (IsSameType<T_IN, float>::value) {251+ if constexpr (IsSameType<T_IN, float>::value) {
252- DataCopy<float, LoadDist::DIST_BRC_B32>(dst, src);252+ LoadAlign<float, LoadDist::DIST_BRC_B32>(dst, src);
253- } else if constexpr (IsSameType<T_IN, int32_t>::value) {253+ } else if constexpr (IsSameType<T_IN, int32_t>::value) {
254- RegTensor<T_IN> xIn;254+ RegTensor<T_IN> xIn;
255- DataCopy<int32_t, LoadDist::DIST_BRC_B32>(xIn, src);255+ LoadAlign<int32_t, LoadDist::DIST_BRC_B32>(xIn, src);
256- Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);256+ Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);
257- } else {257+ } else {
258- RegTensor<T_IN> xIn;258+ RegTensor<T_IN> xIn;
259- DataCopy<T_IN, LoadDist::DIST_BRC_B16>(xIn, src);259+ LoadAlign<T_IN, LoadDist::DIST_BRC_B16>(xIn, src);
260- Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg);260+ Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg);
261- }261+ }
262- }262+ }
263- 263+ 
264- template <bool IsBody>264+ template <bool IsBody>
265- __aicore__ inline void ComputeMul(LocalTensor<float>& rstdLocal, int64_t count)265+ __aicore__ inline void ComputeMul(LocalTensor<float>& rstdLocal, int64_t count)
266- {266+ {
267- LocalTensor<float> reduceLocal = reduceBuf_.Get<float>();267+ LocalTensor<float> reduceLocal = reduceBuf_.Get<float>();
268- LocalTensor<float> gammaLocal = inQueueGamma_.DeQue<float>();268+ LocalTensor<float> gammaLocal = inQueueGamma_.DeQue<float>();
269- LocalTensor<float> dyLocal = inQueueDy_.DeQue<float>();269+ LocalTensor<float> dyLocal = inQueueDy_.DeQue<float>();
270- LocalTensor<float> xLocal = inQueueX_.DeQue<float>();270+ LocalTensor<float> xLocal = inQueueX_.DeQue<float>();
271- 271+ 
272- uint32_t sreg = count;272+ uint32_t sreg = count;
273- constexpr uint32_t oneRepeat = V_LENGTH;273+ constexpr uint32_t oneRepeat = V_LENGTH;
274- uint16_t repeatCount = DivCeil(count, oneRepeat);274+ uint16_t repeatCount = DivCeil(count, oneRepeat);
275- __local_mem__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();275+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
276- __local_mem__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();276+ __ubuf__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();
277- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();277+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
278- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();278+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
279- __local_mem__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();279+ __ubuf__ float* reduceAddr = (__ubuf__ float*)reduceLocal.GetPhyAddr();
280- __VEC_SCOPE__280+ __VEC_SCOPE__
281- {281+ {
282- RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;282+ RegTensor<float> gammaReg, dyReg, xReg, rstdReg, mulReg0, mulReg2, mulReg3;
283- MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();283+ MaskReg maskReg = CreateMask<float, MaskPattern::ALL>();
284- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);284+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);
285- for (uint16_t i = 0; i < repeatCount; i++) {285+ for (uint16_t i = 0; i < repeatCount; i++) {
286- LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);286+ LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
287- LoadAndCast(dyReg, dyAddr, maskReg, i * oneRepeat);287+ LoadAndCast(dyReg, dyAddr, maskReg, i * oneRepeat);
288- Mul(mulReg2, dyReg, gammaReg, maskReg);288+ Mul(mulReg2, dyReg, gammaReg, maskReg);
289- LoadAndCast(xReg, xAddr, maskReg, i * oneRepeat);289+ LoadAndCast(xReg, xAddr, maskReg, i * oneRepeat);
290- Mul(mulReg0, xReg, rstdReg, maskReg);290+ Mul(mulReg0, xReg, rstdReg, maskReg);
291- Mul(mulReg3, mulReg2, mulReg0, maskReg);291+ Mul(mulReg3, mulReg2, mulReg0, maskReg);
292- if constexpr (IsBody) {292+ if constexpr (IsBody) {
293- DataCopy(reduceAddr + static_cast<uint32_t>(i * oneRepeat), mulReg3, maskReg);293+ StoreAlign(reduceAddr + static_cast<uint32_t>(i * oneRepeat), mulReg3, maskReg);
294- } else {294+ } else {
295- DataCopy(reduceAddr + static_cast<uint32_t>(ubFactorD_ + i * oneRepeat), mulReg3,295+ StoreAlign(reduceAddr + static_cast<uint32_t>(ubFactorD_ + i * oneRepeat), mulReg3,
296- maskReg); // 注意补零296+ maskReg); // 注意补零
297- }297+ }
298- }298+ }
299- }299+ }
300- 300+ 
301- inQueueGamma_.FreeTensor(gammaLocal);301+ inQueueGamma_.FreeTensor(gammaLocal);
302- inQueueDy_.FreeTensor(dyLocal);302+ inQueueDy_.FreeTensor(dyLocal);
303- inQueueX_.FreeTensor(xLocal);303+ inQueueX_.FreeTensor(xLocal);
304- }304+ }
305- 305+ 
306- __aicore__ inline void WholeReduceSum(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal, int64_t count,306+ __aicore__ inline void WholeReduceSum(LocalTensor<float>& dstLocal, LocalTensor<float>& srcLocal, int64_t count,
307- int32_t dstOffset)307+ int32_t dstOffset)
308- {308+ {
309- // 对齐到512BYTE, reduce需要309+ // 对齐到512BYTE, reduce需要
310- int64_t countBlockAlign = AlignUp(count, FLOAT_NUM_BLOCK); // 搬入已对齐310+ int64_t countBlockAlign = AlignUp(count, FLOAT_NUM_BLOCK); // 搬入已对齐
311- int64_t count2VLAlign = AlignUp(count, FLOAT_NUM_2VL);311+ int64_t count2VLAlign = AlignUp(count, FLOAT_NUM_2VL);
312- if (count2VLAlign - countBlockAlign > 0) {312+ if (count2VLAlign - countBlockAlign > 0) {
313- Duplicate(srcLocal[countBlockAlign], 0.0f, count2VLAlign - countBlockAlign);313+ Duplicate(srcLocal[countBlockAlign], 0.0f, count2VLAlign - countBlockAlign);
314- }314+ }
315- int64_t power = count2VLAlign < NUM_TWO * ubFactorD_ ?315+ int64_t power = count2VLAlign < NUM_TWO * ubFactorD_ ?
316- ubFactorD_ :316+ ubFactorD_ :
317- NUM_TWO * ubFactorD_; // 等于2*UbFactorD_时设为相同大小，否则为其一半317+ NUM_TWO * ubFactorD_; // 等于2*UbFactorD_时设为相同大小，否则为其一半
318- LocalTensor<float> workLocal = workBuf_.Get<float>();318+ LocalTensor<float> workLocal = workBuf_.Get<float>();
319- ReduceSumImpl(dstLocal, srcLocal, workLocal, dstOffset, count2VLAlign, power);319+ ReduceSumImpl(dstLocal, srcLocal, workLocal, dstOffset, count2VLAlign, power);
320- }320+ }
321- 321+ 
322- __aicore__ inline void LatterProcess(LocalTensor<float>& rstdLocal, int64_t rowIdx)322+ __aicore__ inline void LatterProcess(LocalTensor<float>& rstdLocal, int64_t rowIdx)
323- {323+ {
324- for (int64_t colIdx = 0; colIdx < cols_; colIdx += ubFactorD_) {324+ for (int64_t colIdx = 0; colIdx < cols_; colIdx += ubFactorD_) {
325- int64_t remainCount = cols_ - colIdx;325+ int64_t remainCount = cols_ - colIdx;
326- int64_t calcCount = Min(remainCount, ubFactorD_);326+ int64_t calcCount = Min(remainCount, ubFactorD_);
327- CopyInGamma(colIdx, calcCount);327+ CopyInGamma(colIdx, calcCount);
328- CopyInDy(rowIdx, colIdx, calcCount);328+ CopyInDy(rowIdx, colIdx, calcCount);
329- CopyInX(rowIdx, colIdx, calcCount);329+ CopyInX(rowIdx, colIdx, calcCount);
330- ComputeLatter(rstdLocal, calcCount);330+ ComputeLatter(rstdLocal, calcCount);
331- CopyOutDx(rowIdx, colIdx, calcCount);331+ CopyOutDx(rowIdx, colIdx, calcCount);
332- }332+ }
333- }333+ }
334- 334+ 
335- __aicore__ inline void ComputeLatter(LocalTensor<float>& rstdLocal, int64_t count)335+ __aicore__ inline void ComputeLatter(LocalTensor<float>& rstdLocal, int64_t count)
336- {336+ {
337- LocalTensor<float> gammaLocal = inQueueGamma_.DeQue<float>();337+ LocalTensor<float> gammaLocal = inQueueGamma_.DeQue<float>();
338- LocalTensor<float> dyLocal = inQueueDy_.DeQue<float>();338+ LocalTensor<float> dyLocal = inQueueDy_.DeQue<float>();
339- LocalTensor<float> xLocal = inQueueX_.DeQue<float>();339+ LocalTensor<float> xLocal = inQueueX_.DeQue<float>();
340- LocalTensor<T_DX> dxLocal = outQueueDx_.AllocTensor<T_DX>();340+ LocalTensor<T_DX> dxLocal = outQueueDx_.AllocTensor<T_DX>();
341- LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();341+ LocalTensor<float> tmpSumLocal = tmpSumBuf_.Get<float>();
342- LocalTensor<T_SCALES_X> scalesXLocal;342+ LocalTensor<T_SCALES_X> scalesXLocal;
343- LocalTensor<T_OFFSET_X> offsetXLocal;343+ LocalTensor<T_OFFSET_X> offsetXLocal;
344- 344+ 
345- uint32_t sreg = count;345+ uint32_t sreg = count;
346- constexpr uint32_t oneRepeat = V_LENGTH;346+ constexpr uint32_t oneRepeat = V_LENGTH;
347- uint16_t repeatCount = DivCeil(count, oneRepeat); // 可能会报错347+ uint16_t repeatCount = DivCeil(count, oneRepeat); // 可能会报错
348- __local_mem__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();348+ __ubuf__ T_GAMMA* gammaAddr = (__ubuf__ T_GAMMA*)gammaLocal.GetPhyAddr();
349- __local_mem__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();349+ __ubuf__ T_DY* dyAddr = (__ubuf__ T_DY*)dyLocal.GetPhyAddr();
350- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();350+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
351- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();351+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
352- __local_mem__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();352+ __ubuf__ float* meanAddr = (__ubuf__ float*)tmpSumLocal.GetPhyAddr();
353- __local_mem__ T_DX* dxAddr = (__ubuf__ T_DX*)dxLocal.GetPhyAddr();353+ __ubuf__ T_DX* dxAddr = (__ubuf__ T_DX*)dxLocal.GetPhyAddr();
354- __local_mem__ T_SCALES_X* scalesXAddr;354+ __ubuf__ T_SCALES_X* scalesXAddr;
355- __local_mem__ T_OFFSET_X* offsetXAddr;355+ __ubuf__ T_OFFSET_X* offsetXAddr;
356- 356+ 
357- scalesXLocal = scalesXBuf_.Get<T_SCALES_X>();357+ scalesXLocal = scalesXBuf_.Get<T_SCALES_X>();
358- scalesXAddr = (__ubuf__ T_SCALES_X*)scalesXLocal.GetPhyAddr();358+ scalesXAddr = (__ubuf__ T_SCALES_X*)scalesXLocal.GetPhyAddr();
359- if constexpr (HAS_OFFSET_X) {359+ if constexpr (HAS_OFFSET_X) {
360- offsetXLocal = offsetXBuf_.Get<T_OFFSET_X>();360+ offsetXLocal = offsetXBuf_.Get<T_OFFSET_X>();
361- offsetXAddr = (__ubuf__ T_OFFSET_X*)offsetXLocal.GetPhyAddr();361+ offsetXAddr = (__ubuf__ T_OFFSET_X*)offsetXLocal.GetPhyAddr();
362- }362+ }
363- 363+ 
364- __VEC_SCOPE__364+ __VEC_SCOPE__
365- {365+ {
366- RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;366+ RegTensor<float> gammaReg, dyReg, xReg, rstdReg, meanReg, dxReg, mulReg0, mulReg2, mulReg4, subReg;
367- RegTensor<float> scalesXReg, scalesXResultReg, offsetXReg;367+ RegTensor<float> scalesXReg, scalesXResultReg, offsetXReg;
368- MaskReg maskReg;368+ MaskReg maskReg;
369- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);369+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr);
370- DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr);370+ LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr);
371- for (uint16_t i = 0; i < repeatCount; i++) {371+ for (uint16_t i = 0; i < repeatCount; i++) {
372- maskReg = UpdateMask<float>(sreg);372+ maskReg = UpdateMask<float>(sreg);
373- LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);373+ LoadAndCast(gammaReg, gammaAddr, maskReg, i * oneRepeat);
374- LoadAndCast(dyReg, dyAddr, maskReg, i * oneRepeat);374+ LoadAndCast(dyReg, dyAddr, maskReg, i * oneRepeat);
375- Mul(mulReg2, dyReg, gammaReg, maskReg);375+ Mul(mulReg2, dyReg, gammaReg, maskReg);
376- LoadAndCast(xReg, xAddr, maskReg, i * oneRepeat);376+ LoadAndCast(xReg, xAddr, maskReg, i * oneRepeat);
377- Mul(mulReg0, xReg, rstdReg, maskReg);377+ Mul(mulReg0, xReg, rstdReg, maskReg);
378- Mul(mulReg4, mulReg0, meanReg, maskReg);378+ Mul(mulReg4, mulReg0, meanReg, maskReg);
379- Sub(subReg, mulReg2, mulReg4, maskReg);379+ Sub(subReg, mulReg2, mulReg4, maskReg);
380- Mul(dxReg, subReg, rstdReg, maskReg);380+ Mul(dxReg, subReg, rstdReg, maskReg);
381- // cal quant381+ // cal quant
382- LoadTensorForDtypeTIn(scalesXAddr, scalesXReg, maskReg);382+ LoadTensorForDtypeTIn(scalesXAddr, scalesXReg, maskReg);
383- if constexpr (DIV_MODE) {383+ if constexpr (DIV_MODE) {
384- Div(scalesXResultReg, dxReg, scalesXReg, maskReg);384+ Div(scalesXResultReg, dxReg, scalesXReg, maskReg);
385- } else {385+ } else {
386- Mul(scalesXResultReg, dxReg, scalesXReg, maskReg);386+ Mul(scalesXResultReg, dxReg, scalesXReg, maskReg);
387- }387+ }
388- if constexpr (HAS_OFFSET_X) {388+ if constexpr (HAS_OFFSET_X) {
389- LoadTensorForDtypeTIn(offsetXAddr, offsetXReg, maskReg);389+ LoadTensorForDtypeTIn(offsetXAddr, offsetXReg, maskReg);
390- Add(scalesXResultReg, scalesXResultReg, offsetXReg, maskReg);390+ Add(scalesXResultReg, scalesXResultReg, offsetXReg, maskReg);
391- }391+ }
392- if constexpr (IsSameType<T_DX, hifloat8_t>::value) {392+ if constexpr (IsSameType<T_DX, hifloat8_t>::value) {
393- RegTensor<T_DX> dxRegHif8;393+ RegTensor<T_DX> dxRegHif8;
394- Cast<T_DX, float, castTraitFp322Hifp8>(dxRegHif8, scalesXResultReg, maskReg);394+ Cast<T_DX, float, castTraitFp322Hifp8>(dxRegHif8, scalesXResultReg, maskReg);
395- DataCopy<T_DX, StoreDist::DIST_PACK4_B32>(dxAddr + static_cast<uint32_t>(i * oneRepeat), dxRegHif8,395+ StoreAlign<T_DX, StoreDist::DIST_PACK4_B32>(dxAddr + static_cast<uint32_t>(i * oneRepeat),
396- maskReg);396+ dxRegHif8, maskReg);
397- } else if constexpr (IsSameType<T_DX, int8_t>::value) {397+ } else if constexpr (IsSameType<T_DX, int8_t>::value) {
398- RegTensor<T_DX> dxRegInt8;398+ RegTensor<T_DX> dxRegInt8;
399- RegTensor<half> dxRegFp16;399+ RegTensor<half> dxRegFp16;
400- RegTensor<int32_t> dxRegInt32;400+ RegTensor<int32_t> dxRegInt32;
401- Cast<int32_t, float, castTraitFp322Int32>(dxRegInt32, scalesXResultReg, maskReg);401+ Cast<int32_t, float, castTraitFp322Int32>(dxRegInt32, scalesXResultReg, maskReg);
402- Cast<float, int32_t, castTraitInt322Fp32>(scalesXResultReg, dxRegInt32, maskReg);402+ Cast<float, int32_t, castTraitInt322Fp32>(scalesXResultReg, dxRegInt32, maskReg);
403- Cast<half, float, castTraitFp322Fp16>(dxRegFp16, scalesXResultReg, maskReg);403+ Cast<half, float, castTraitFp322Fp16>(dxRegFp16, scalesXResultReg, maskReg);
404- Cast<T_DX, half, castTraitFp162Int8>(dxRegInt8, dxRegFp16, maskReg);404+ Cast<T_DX, half, castTraitFp162Int8>(dxRegInt8, dxRegFp16, maskReg);
405- DataCopy<T_DX, StoreDist::DIST_PACK4_B32>(dxAddr + static_cast<uint32_t>(i * oneRepeat), dxRegInt8,405+ StoreAlign<T_DX, StoreDist::DIST_PACK4_B32>(dxAddr + static_cast<uint32_t>(i * oneRepeat),
406- maskReg);406+ dxRegInt8, maskReg);
407- }407+ }
408- }408+ }
409- }409+ }
410- 410+ 
411- inQueueGamma_.FreeTensor(gammaLocal);411+ inQueueGamma_.FreeTensor(gammaLocal);
412- inQueueX_.FreeTensor(xLocal);412+ inQueueX_.FreeTensor(xLocal);
413- inQueueDy_.FreeTensor(dyLocal);413+ inQueueDy_.FreeTensor(dyLocal);
414- outQueueDx_.EnQue(dxLocal);414+ outQueueDx_.EnQue(dxLocal);
415- }415+ }
416- 416+ 
417- __aicore__ inline void CopyOutDx(int64_t rowIdx, int64_t colIdx, int64_t count)417+ __aicore__ inline void CopyOutDx(int64_t rowIdx, int64_t colIdx, int64_t count)
418- {418+ {
419- LocalTensor<T_DX> dxLocal = outQueueDx_.DeQue<T_DX>();419+ LocalTensor<T_DX> dxLocal = outQueueDx_.DeQue<T_DX>();
420- DataCopyExtParams copyParams{420+ DataCopyExtParams copyParams{
421- 1, // blockCount421+ 1, // blockCount
422- static_cast<uint32_t>(count * sizeof(T_DX)), // blockLen422+ static_cast<uint32_t>(count * sizeof(T_DX)), // blockLen
423- 0, // srcStride423+ 0, // srcStride
424- 0, // dstStride424+ 0, // dstStride
425- 0 // rsv425+ 0 // rsv
426- };426+ };
427- DataCopyPad(dxGm_[rowIdx * cols_ + colIdx], dxLocal, copyParams);427+ DataCopyPad(dxGm_[rowIdx * cols_ + colIdx], dxLocal, copyParams);
428- outQueueDx_.FreeTensor(dxLocal);428+ outQueueDx_.FreeTensor(dxLocal);
429- }429+ }
430- 430+ 
431-private:431+private:
432- TPipe* Ppipe_;432+ TPipe* Ppipe_;
433- const RmsNormGradQuantRegbaseDxTilingData* tiling_;433+ const RmsNormGradQuantRegbaseDxTilingData* tiling_;
434- GlobalTensor<T_DY> dyGm_;434+ GlobalTensor<T_DY> dyGm_;
435- GlobalTensor<T_X> xGm_;435+ GlobalTensor<T_X> xGm_;
436- GlobalTensor<T_GAMMA> gammaGm_;436+ GlobalTensor<T_GAMMA> gammaGm_;
437- GlobalTensor<float> rstdGm_;437+ GlobalTensor<float> rstdGm_;
438- GlobalTensor<T_DX> dxGm_;438+ GlobalTensor<T_DX> dxGm_;
439- GlobalTensor<T_SCALES_X> scalesXGm_;439+ GlobalTensor<T_SCALES_X> scalesXGm_;
440- GlobalTensor<T_OFFSET_X> offsetXGm_;440+ GlobalTensor<T_OFFSET_X> offsetXGm_;
441- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueDy_;441+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueDy_;
442- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueX_;442+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueX_;
443- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueRstd_;443+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueRstd_;
444- TQue<QuePosition::VECIN, DEPTH_TWO> inQueueGamma_;444+ TQue<QuePosition::VECIN, DEPTH_TWO> inQueueGamma_;
445- TQue<QuePosition::VECOUT, DEPTH_TWO> outQueueDx_;445+ TQue<QuePosition::VECOUT, DEPTH_TWO> outQueueDx_;
446- TBuf<TPosition::VECCALC> reduceBuf_;446+ TBuf<TPosition::VECCALC> reduceBuf_;
447- TBuf<TPosition::VECCALC> level0Buf_;447+ TBuf<TPosition::VECCALC> level0Buf_;
448- TBuf<TPosition::VECCALC> level1Buf_;448+ TBuf<TPosition::VECCALC> level1Buf_;
449- TBuf<TPosition::VECCALC> level2Buf_;449+ TBuf<TPosition::VECCALC> level2Buf_;
450- TBuf<TPosition::VECCALC> tmpSumBuf_;450+ TBuf<TPosition::VECCALC> tmpSumBuf_;
451- TBuf<TPosition::VECCALC> workBuf_;451+ TBuf<TPosition::VECCALC> workBuf_;
452- TBuf<TPosition::VECCALC> scalesXBuf_;452+ TBuf<TPosition::VECCALC> scalesXBuf_;
453- TBuf<TPosition::VECCALC> offsetXBuf_;453+ TBuf<TPosition::VECCALC> offsetXBuf_;
454- uint32_t usedCoreNum_;454+ uint32_t usedCoreNum_;
455- int64_t rows_;455+ int64_t rows_;
456- int64_t cols_;456+ int64_t cols_;
457- int64_t blockFactor_;457+ int64_t blockFactor_;
458- int64_t ubFactorD_;458+ int64_t ubFactorD_;
459- int64_t bodyPart_;459+ int64_t bodyPart_;
460- float avgFactor1_;460+ float avgFactor1_;
461-};461+};
462-} // namespace RmsNormGradQuant462+} // namespace RmsNormGradQuant
463-#endif // RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H463+#endif // RMS_NORM_GRAD_REGBASE_DX_SPLIT_D_H
Mnorm/rms_norm_quant_v2/op_kernel/arch35/rms_norm_quant_v2_regbase_common.h+15-17
@@ -118,53 +118,51 @@ __aicore__ inline int64_t GetCacheId(const int64_t idx)
118}118}
119 119 
120template <typename T_IN>120template <typename T_IN>
121-__aicore__ inline void LoadScalarForDtypeTIn(__local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg,121+__aicore__ inline void LoadScalarForDtypeTIn(__ubuf__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)
122- uint32_t offset)
123{122{
124 if constexpr (IsSameType<T_IN, float>::value) {123 if constexpr (IsSameType<T_IN, float>::value) {
125- DataCopy<float, LoadDist::DIST_BRC_B32>(dst, src + offset);124+ LoadAlign<float, LoadDist::DIST_BRC_B32>(dst, src + offset);
126 } else if constexpr (IsSameType<T_IN, int8_t>::value) {125 } else if constexpr (IsSameType<T_IN, int8_t>::value) {
127 RegTensor<T_IN> xIn;126 RegTensor<T_IN> xIn;
128 RegTensor<half> xInTmp;127 RegTensor<half> xInTmp;
129- DataCopy<T_IN, LoadDist::DIST_BRC_B8>(xIn, src + offset);128+ LoadAlign<T_IN, LoadDist::DIST_BRC_B8>(xIn, src + offset);
130 Cast<half, T_IN, castTraitI82F16>(xInTmp, xIn, preg);129 Cast<half, T_IN, castTraitI82F16>(xInTmp, xIn, preg);
131 Cast<float, half, castTraitF162F32>(dst, xInTmp, preg);130 Cast<float, half, castTraitF162F32>(dst, xInTmp, preg);
132 } else if constexpr (IsSameType<T_IN, int32_t>::value) {131 } else if constexpr (IsSameType<T_IN, int32_t>::value) {
133 RegTensor<T_IN> xIn;132 RegTensor<T_IN> xIn;
134- DataCopy<int32_t, LoadDist::DIST_BRC_B32>(xIn, src + offset);133+ LoadAlign<int32_t, LoadDist::DIST_BRC_B32>(xIn, src + offset);
135 Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);134 Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);
136 } else {135 } else {
137 RegTensor<T_IN> xIn;136 RegTensor<T_IN> xIn;
138- DataCopy<T_IN, LoadDist::DIST_BRC_B16>(xIn, src + offset);137+ LoadAlign<T_IN, LoadDist::DIST_BRC_B16>(xIn, src + offset);
139 Cast<float, T_IN, castTraitF162F32>(dst, xIn, preg);138 Cast<float, T_IN, castTraitF162F32>(dst, xIn, preg);
140 }139 }
141}140}
142 141 
143template <typename T_IN>142template <typename T_IN>
144-__aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg,143+__aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset)
145- uint32_t offset)
146{144{
147 if constexpr (IsSameType<T_IN, float>::value) {145 if constexpr (IsSameType<T_IN, float>::value) {
148- DataCopy<float, LoadDist::DIST_NORM>(dst, src + offset);146+ LoadAlign<float, LoadDist::DIST_NORM>(dst, src + offset);
149 } else if constexpr (IsSameType<T_IN, int8_t>::value) {147 } else if constexpr (IsSameType<T_IN, int8_t>::value) {
150 RegTensor<T_IN> xIn;148 RegTensor<T_IN> xIn;
151 RegTensor<half> xInTmp;149 RegTensor<half> xInTmp;
152- DataCopy<T_IN, LoadDist::DIST_UNPACK4_B8>(xIn, src + offset);150+ LoadAlign<T_IN, LoadDist::DIST_UNPACK4_B8>(xIn, src + offset);
153 Cast<half, T_IN, castTraitI82F16>(xInTmp, xIn, preg);151 Cast<half, T_IN, castTraitI82F16>(xInTmp, xIn, preg);
154 Cast<float, half, castTraitF162F32>(dst, xInTmp, preg);152 Cast<float, half, castTraitF162F32>(dst, xInTmp, preg);
155 } else if constexpr (IsSameType<T_IN, int32_t>::value) {153 } else if constexpr (IsSameType<T_IN, int32_t>::value) {
156 RegTensor<T_IN> xIn;154 RegTensor<T_IN> xIn;
157- DataCopy<int32_t, LoadDist::DIST_NORM>(xIn, src + offset);155+ LoadAlign<int32_t, LoadDist::DIST_NORM>(xIn, src + offset);
158 Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);156 Cast<float, T_IN, castTraitInt322Fp32>(dst, xIn, preg);
159 } else {157 } else {
160 RegTensor<T_IN> xIn;158 RegTensor<T_IN> xIn;
161- DataCopy<T_IN, LoadDist::DIST_UNPACK_B16>(xIn, src + offset);159+ LoadAlign<T_IN, LoadDist::DIST_UNPACK_B16>(xIn, src + offset);
162 Cast<float, T_IN, castTraitF162F32>(dst, xIn, preg);160 Cast<float, T_IN, castTraitF162F32>(dst, xIn, preg);
163 }161 }
164}162}
165 163 
166template <typename T_OUT>164template <typename T_OUT>
167-__aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_OUT* dst, RegTensor<float>& xRegFp32, MaskReg& preg,165+__aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T_OUT* dst, RegTensor<float>& xRegFp32, MaskReg& preg,
168 MaskReg& mask4Int4, uint32_t offset)166 MaskReg& mask4Int4, uint32_t offset)
169{167{
170 if constexpr (IsSameType<T_OUT, int8_t>::value) {168 if constexpr (IsSameType<T_OUT, int8_t>::value) {
@@ -172,7 +170,7 @@ __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_OUT* dst, RegTens
172 RegTensor<half> xRegFp16;170 RegTensor<half> xRegFp16;
173 Cast<half, float, castTraitFp322Fp16>(xRegFp16, xRegFp32, preg);171 Cast<half, float, castTraitFp322Fp16>(xRegFp16, xRegFp32, preg);
174 Cast<int8_t, half, castTraitFp162Int8>(xOut, xRegFp16, preg);172 Cast<int8_t, half, castTraitFp162Int8>(xOut, xRegFp16, preg);
175- DataCopy<T_OUT, StoreDist::DIST_PACK4_B32>(dst + offset, xOut, preg);173+ StoreAlign<T_OUT, StoreDist::DIST_PACK4_B32>(dst + offset, xOut, preg);
176 } else if constexpr (IsSameType<T_OUT, uint8_t>::value) {174 } else if constexpr (IsSameType<T_OUT, uint8_t>::value) {
177 RegTensor<half> xRegFp16;175 RegTensor<half> xRegFp16;
178 RegTensor<uint16_t> xUint16;176 RegTensor<uint16_t> xUint16;
@@ -180,15 +178,15 @@ __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_OUT* dst, RegTens
180 Cast<half, float, castTraitFp322Fp16>(xRegFp16, xRegFp32, preg);178 Cast<half, float, castTraitFp322Fp16>(xRegFp16, xRegFp32, preg);
181 Pack(xUint16, (RegTensor<uint32_t>&)xRegFp16);179 Pack(xUint16, (RegTensor<uint32_t>&)xRegFp16);
182 Cast<int4x2_t, half, castTraitFp162Int8>((RegTensor<int4x2_t>&)xUint8, (RegTensor<half>&)xUint16, preg);180 Cast<int4x2_t, half, castTraitFp162Int8>((RegTensor<int4x2_t>&)xUint8, (RegTensor<half>&)xUint16, preg);
183- DataCopy<uint8_t, StoreDist::DIST_PACK4_B32>(dst + offset / 2, xUint8, mask4Int4);181+ StoreAlign<uint8_t, StoreDist::DIST_PACK4_B32>(dst + offset / 2, xUint8, mask4Int4);
184 } else if constexpr (IsSameType<T_OUT, fp8_e4m3fn_t>::value || IsSameType<T_OUT, fp8_e5m2_t>::value) {182 } else if constexpr (IsSameType<T_OUT, fp8_e4m3fn_t>::value || IsSameType<T_OUT, fp8_e5m2_t>::value) {
185 RegTensor<T_OUT> xOut;183 RegTensor<T_OUT> xOut;
186 Cast<T_OUT, float, castTraitFp322Fp8>(xOut, xRegFp32, preg);184 Cast<T_OUT, float, castTraitFp322Fp8>(xOut, xRegFp32, preg);
187- DataCopy<T_OUT, StoreDist::DIST_PACK4_B32>(dst + offset, xOut, preg);185+ StoreAlign<T_OUT, StoreDist::DIST_PACK4_B32>(dst + offset, xOut, preg);
188 } else if constexpr (IsSameType<T_OUT, hifloat8_t>::value) {186 } else if constexpr (IsSameType<T_OUT, hifloat8_t>::value) {
189 RegTensor<T_OUT> xOut;187 RegTensor<T_OUT> xOut;
190 Cast<T_OUT, float, castTraitFp322Hifp8>(xOut, xRegFp32, preg);188 Cast<T_OUT, float, castTraitFp322Hifp8>(xOut, xRegFp32, preg);
191- DataCopy<T_OUT, StoreDist::DIST_PACK4_B32>(dst + offset, xOut, preg);189+ StoreAlign<T_OUT, StoreDist::DIST_PACK4_B32>(dst + offset, xOut, preg);
192 }190 }
193}191}
194} // namespace RmsNormQuantV2192} // namespace RmsNormQuantV2
Mnorm/rms_norm_quant_v2/op_kernel/arch35/rms_norm_quant_v2_regbase_full_load.h+775-775
@@ -1,775 +1,775 @@
1-/**1+/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
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-/*!11+/*!
12- * \file rms_norm_quant_v2_regbase_full_load.h12+ * \file rms_norm_quant_v2_regbase_full_load.h
13- * \brief13+ * \brief
14- */14+ */
15-#ifndef RMS_NORM_QUANT_V2_REBASE_FULL_LOAD_H_15+#ifndef RMS_NORM_QUANT_V2_REBASE_FULL_LOAD_H_
16-#define RMS_NORM_QUANT_V2_REBASE_FULL_LOAD_H_16+#define RMS_NORM_QUANT_V2_REBASE_FULL_LOAD_H_
17-#include "kernel_utils.h"17+#include "kernel_utils.h"
18-#include "../inc/platform.h"18+#include "../inc/platform.h"
19-#include "rms_norm_quant_v2_regbase_common.h"19+#include "rms_norm_quant_v2_regbase_common.h"
20-#include "../../norm_common/reduce_common_regbase.h"20+#include "../../norm_common/reduce_common_regbase.h"
21- 21+ 
22-namespace RmsNormQuantV2 {22+namespace RmsNormQuantV2 {
23- 23+ 
24-template <typename T_X, typename T_Y, typename T_SCALES, typename T_ZEROPOINTS>24+template <typename T_X, typename T_Y, typename T_SCALES, typename T_ZEROPOINTS>
25-class RmsNormQuantV2RegbaseFullLoad {25+class RmsNormQuantV2RegbaseFullLoad {
26-private:26+private:
27- TPipe* pipe_ = nullptr;27+ TPipe* pipe_ = nullptr;
28- using yDtype = std::conditional_t<IsSameType<T_Y, int4b_t>::value, uint8_t, T_Y>;28+ using yDtype = std::conditional_t<IsSameType<T_Y, int4b_t>::value, uint8_t, T_Y>;
29- // GM Buffer29+ // GM Buffer
30- GlobalTensor<T_X> xGm, gammaGm, betaGm;30+ GlobalTensor<T_X> xGm, gammaGm, betaGm;
31- GlobalTensor<T_SCALES> scales1Gm, scales2Gm;31+ GlobalTensor<T_SCALES> scales1Gm, scales2Gm;
32- GlobalTensor<T_ZEROPOINTS> zeroPoints1Gm, zeroPoints2Gm;32+ GlobalTensor<T_ZEROPOINTS> zeroPoints1Gm, zeroPoints2Gm;
33- GlobalTensor<yDtype> y1Gm, y2Gm;33+ GlobalTensor<yDtype> y1Gm, y2Gm;
34- GlobalTensor<float> rstdGm;34+ GlobalTensor<float> rstdGm;
35- // UB Buffer35+ // UB Buffer
36- TQue<QuePosition::VECIN, 1> inQueueX;36+ TQue<QuePosition::VECIN, 1> inQueueX;
37- // gamma beta scales0 scales1 zero_points0 zero_point1 all in this queue37+ // gamma beta scales0 scales1 zero_points0 zero_point1 all in this queue
38- TQue<QuePosition::VECIN, 1> inQueueOhter;38+ TQue<QuePosition::VECIN, 1> inQueueOhter;
39- TQue<QuePosition::VECOUT, 1> outQueueY1, outQueueY2;39+ TQue<QuePosition::VECOUT, 1> outQueueY1, outQueueY2;
40- TQue<QuePosition::VECOUT, 1> outQueueRstd;40+ TQue<QuePosition::VECOUT, 1> outQueueRstd;
41- TBuf<TPosition::VECCALC> rstdBuf;41+ TBuf<TPosition::VECCALC> rstdBuf;
42- TBuf<TPosition::VECCALC> reduceTmpBuf;42+ TBuf<TPosition::VECCALC> reduceTmpBuf;
43- 43+ 
44- LocalTensor<T_X> gammaLocal;44+ LocalTensor<T_X> gammaLocal;
45- LocalTensor<T_SCALES> scales1Local, scales2Local;45+ LocalTensor<T_SCALES> scales1Local, scales2Local;
46- LocalTensor<T_ZEROPOINTS> zeroPoints1Local, zeroPoints2Local;46+ LocalTensor<T_ZEROPOINTS> zeroPoints1Local, zeroPoints2Local;
47- LocalTensor<T_X> betaLocal;47+ LocalTensor<T_X> betaLocal;
48- 48+ 
49- // Tiling data49+ // Tiling data
50- int64_t numA{0};50+ int64_t numA{0};
51- int64_t numR{0};51+ int64_t numR{0};
52- int64_t numQ{0};52+ int64_t numQ{0};
53- int64_t blockFactor{0};53+ int64_t blockFactor{0};
54- int64_t blockTail{0};54+ int64_t blockTail{0};
55- int64_t ubFactor{0};55+ int64_t ubFactor{0};
56- int64_t binaryAdd{0};56+ int64_t binaryAdd{0};
57- uint32_t optionMask{0};57+ uint32_t optionMask{0};
58- bool isScaleDiv{0};58+ bool isScaleDiv{0};
59- float epsilon{0};59+ float epsilon{0};
60- float avgFactor{0};60+ float avgFactor{0};
61- // Platform61+ // Platform
62- int64_t blockIdx{0};62+ int64_t blockIdx{0};
63- int64_t blockNum{0};63+ int64_t blockNum{0};
64- int64_t oriOverflowMode{0};64+ int64_t oriOverflowMode{0};
65- uint32_t blockSize = platform::GetUbBlockSize();65+ uint32_t blockSize = platform::GetUbBlockSize();
66- uint32_t vectorLen = platform::GetVRegSize();66+ uint32_t vectorLen = platform::GetVRegSize();
67- uint32_t blockSizeB32 = platform::GetUbBlockSize() / sizeof(float);67+ uint32_t blockSizeB32 = platform::GetUbBlockSize() / sizeof(float);
68- uint32_t vectorLenB32 = platform::GetVRegSize() / sizeof(float);68+ uint32_t vectorLenB32 = platform::GetVRegSize() / sizeof(float);
69- 69+ 
70- // dtypeSize70+ // dtypeSize
71- uint32_t xDtypeSize{1};71+ uint32_t xDtypeSize{1};
72- uint32_t scalesDtypeSize{1};72+ uint32_t scalesDtypeSize{1};
73- uint32_t zeroPointsDtypeSize{1};73+ uint32_t zeroPointsDtypeSize{1};
74- uint32_t yDtypeSize{1};74+ uint32_t yDtypeSize{1};
75- 75+ 
76- // align value76+ // align value
77- int64_t xGammaBetaAlign{32};77+ int64_t xGammaBetaAlign{32};
78- int64_t scalesAlign{32};78+ int64_t scalesAlign{32};
79- int64_t zeroPointsAlign{32};79+ int64_t zeroPointsAlign{32};
80- int64_t yAlign{32};80+ int64_t yAlign{32};
81- int64_t rstdAlign{32};81+ int64_t rstdAlign{32};
82- 82+ 
83- // calculate value83+ // calculate value
84- int64_t curBlockFactor{0};84+ int64_t curBlockFactor{0};
85- int64_t curUbLoops{0};85+ int64_t curUbLoops{0};
86- int64_t ubFactorTail{0};86+ int64_t ubFactorTail{0};
87- 87+ 
88- // option value88+ // option value
89- bool hasZeroPoints1{false};89+ bool hasZeroPoints1{false};
90- bool hasScales2{false};90+ bool hasScales2{false};
91- bool hasZeroPoints2{false};91+ bool hasZeroPoints2{false};
92- bool hasBeta{false};92+ bool hasBeta{false};
93- bool hasY2{false};93+ bool hasY2{false};
94- uint32_t rstdFlag_{0};94+ uint32_t rstdFlag_{0};
95- 95+ 
96- // option mask const value96+ // option mask const value
97- static constexpr uint32_t SCALES2_MASK = 0b0001;97+ static constexpr uint32_t SCALES2_MASK = 0b0001;
98- static constexpr uint32_t ZEROS_POINTS1_MASK = 0b0010;98+ static constexpr uint32_t ZEROS_POINTS1_MASK = 0b0010;
99- static constexpr uint32_t ZEROS_POINTS2_MASK = 0b0100;99+ static constexpr uint32_t ZEROS_POINTS2_MASK = 0b0100;
100- static constexpr uint32_t BETA_MASK = 0b1000;100+ static constexpr uint32_t BETA_MASK = 0b1000;
101- static constexpr uint32_t DOUBLE_BUFFER_NUM = 2;101+ static constexpr uint32_t DOUBLE_BUFFER_NUM = 2;
102- 102+ 
103- static constexpr float RMS_POS_INF = 3.40282366920938E+38;103+ static constexpr float RMS_POS_INF = 3.40282366920938E+38;
104- static constexpr float RMS_ZERO = 0.0f;104+ static constexpr float RMS_ZERO = 0.0f;
105- 105+ 
106-public:106+public:
107- __aicore__ inline RmsNormQuantV2RegbaseFullLoad(TPipe* pipe) { pipe_ = pipe; }107+ __aicore__ inline RmsNormQuantV2RegbaseFullLoad(TPipe* pipe) { pipe_ = pipe; }
108- 108+ 
109- __aicore__ inline void Init(GM_ADDR x, GM_ADDR gamma, GM_ADDR scales1, GM_ADDR scales2, GM_ADDR zeroPoints1,109+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR gamma, GM_ADDR scales1, GM_ADDR scales2, GM_ADDR zeroPoints1,
110- GM_ADDR zeroPoints2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, GM_ADDR rstd,110+ GM_ADDR zeroPoints2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, GM_ADDR rstd,
111- const RmsNormQuantV2RegbaseFullLoadTilingData* tilingData)111+ const RmsNormQuantV2RegbaseFullLoadTilingData* tilingData)
112- {112+ {
113- // Tiling data113+ // Tiling data
114- numA = tilingData->a;114+ numA = tilingData->a;
115- numR = tilingData->r;115+ numR = tilingData->r;
116- numQ = tilingData->q;116+ numQ = tilingData->q;
117- blockFactor = tilingData->blockFactor;117+ blockFactor = tilingData->blockFactor;
118- blockTail = tilingData->blockTail;118+ blockTail = tilingData->blockTail;
119- ubFactor = tilingData->ubFactor;119+ ubFactor = tilingData->ubFactor;
120- binaryAdd = tilingData->binaryAdd;120+ binaryAdd = tilingData->binaryAdd;
121- optionMask = tilingData->optionMask & 0xF;121+ optionMask = tilingData->optionMask & 0xF;
122- isScaleDiv = tilingData->divMode == 1;122+ isScaleDiv = tilingData->divMode == 1;
123- epsilon = tilingData->epsilon;123+ epsilon = tilingData->epsilon;
124- avgFactor = tilingData->avgFactor;124+ avgFactor = tilingData->avgFactor;
125- rstdFlag_ = tilingData->rstdFlag;125+ rstdFlag_ = tilingData->rstdFlag;
126- 126+ 
127- // dtype size127+ // dtype size
128- xDtypeSize = blockSize / sizeof(T_X);128+ xDtypeSize = blockSize / sizeof(T_X);
129- scalesDtypeSize = blockSize / sizeof(T_SCALES);129+ scalesDtypeSize = blockSize / sizeof(T_SCALES);
130- zeroPointsDtypeSize = blockSize / sizeof(T_ZEROPOINTS);130+ zeroPointsDtypeSize = blockSize / sizeof(T_ZEROPOINTS);
131- yDtypeSize = blockSize / sizeof(yDtype);131+ yDtypeSize = blockSize / sizeof(yDtype);
132- if constexpr (IsSameType<T_Y, int4b_t>::value) {132+ if constexpr (IsSameType<T_Y, int4b_t>::value) {
133- yDtypeSize = yDtypeSize * 2;133+ yDtypeSize = yDtypeSize * 2;
134- }134+ }
135- 135+ 
136- // dtype align136+ // dtype align
137- xGammaBetaAlign = CeilDiv(numR, static_cast<int64_t>(xDtypeSize)) * static_cast<int64_t>(xDtypeSize);137+ xGammaBetaAlign = CeilDiv(numR, static_cast<int64_t>(xDtypeSize)) * static_cast<int64_t>(xDtypeSize);
138- scalesAlign = CeilDiv(numR, static_cast<int64_t>(scalesDtypeSize)) * static_cast<int64_t>(scalesDtypeSize);138+ scalesAlign = CeilDiv(numR, static_cast<int64_t>(scalesDtypeSize)) * static_cast<int64_t>(scalesDtypeSize);
139- zeroPointsAlign = CeilDiv(numR, static_cast<int64_t>(zeroPointsDtypeSize)) *139+ zeroPointsAlign = CeilDiv(numR, static_cast<int64_t>(zeroPointsDtypeSize)) *
140- static_cast<int64_t>(zeroPointsDtypeSize);140+ static_cast<int64_t>(zeroPointsDtypeSize);
141- yAlign = CeilDiv(numR, static_cast<int64_t>(yDtypeSize)) * static_cast<int64_t>(yDtypeSize);141+ yAlign = CeilDiv(numR, static_cast<int64_t>(yDtypeSize)) * static_cast<int64_t>(yDtypeSize);
142- rstdAlign = CeilDiv(ubFactor, static_cast<int64_t>(blockSizeB32)) * static_cast<int64_t>(blockSizeB32);142+ rstdAlign = CeilDiv(ubFactor, static_cast<int64_t>(blockSizeB32)) * static_cast<int64_t>(blockSizeB32);
143- 143+ 
144- blockNum = GetBlockNum();144+ blockNum = GetBlockNum();
145- blockIdx = GetBlockIdx();145+ blockIdx = GetBlockIdx();
146- oriOverflowMode = GetOverflowMode<T_Y>();146+ oriOverflowMode = GetOverflowMode<T_Y>();
147- 147+ 
148- // init option148+ // init option
149- if ((optionMask & ZEROS_POINTS1_MASK) == ZEROS_POINTS1_MASK) {149+ if ((optionMask & ZEROS_POINTS1_MASK) == ZEROS_POINTS1_MASK) {
150- hasZeroPoints1 = true;150+ hasZeroPoints1 = true;
151- }151+ }
152- if ((optionMask & SCALES2_MASK) == SCALES2_MASK) {152+ if ((optionMask & SCALES2_MASK) == SCALES2_MASK) {
153- hasScales2 = true;153+ hasScales2 = true;
154- }154+ }
155- if ((optionMask & ZEROS_POINTS2_MASK) == ZEROS_POINTS2_MASK) {155+ if ((optionMask & ZEROS_POINTS2_MASK) == ZEROS_POINTS2_MASK) {
156- hasZeroPoints2 = true;156+ hasZeroPoints2 = true;
157- }157+ }
158- if ((optionMask & BETA_MASK) == BETA_MASK) {158+ if ((optionMask & BETA_MASK) == BETA_MASK) {
159- hasBeta = true;159+ hasBeta = true;
160- }160+ }
161- hasY2 = hasScales2;161+ hasY2 = hasScales2;
162- 162+ 
163- // init curBlockFactor163+ // init curBlockFactor
164- curBlockFactor = blockIdx == (blockNum - 1) ? blockTail : blockFactor;164+ curBlockFactor = blockIdx == (blockNum - 1) ? blockTail : blockFactor;
165- curUbLoops = CeilDiv(curBlockFactor, ubFactor);165+ curUbLoops = CeilDiv(curBlockFactor, ubFactor);
166- ubFactorTail = curBlockFactor - (curUbLoops - 1) * ubFactor;166+ ubFactorTail = curBlockFactor - (curUbLoops - 1) * ubFactor;
167- 167+ 
168- InitBuffer(x, gamma, scales1, scales2, zeroPoints1, zeroPoints2, beta, y1, y2, rstd);168+ InitBuffer(x, gamma, scales1, scales2, zeroPoints1, zeroPoints2, beta, y1, y2, rstd);
169- }169+ }
170- 170+ 
171- __aicore__ inline void InitBuffer(GM_ADDR x, GM_ADDR gamma, GM_ADDR scales1, GM_ADDR scales2, GM_ADDR zeroPoints1,171+ __aicore__ inline void InitBuffer(GM_ADDR x, GM_ADDR gamma, GM_ADDR scales1, GM_ADDR scales2, GM_ADDR zeroPoints1,
172- GM_ADDR zeroPoints2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, GM_ADDR rstd)172+ GM_ADDR zeroPoints2, GM_ADDR beta, GM_ADDR y1, GM_ADDR y2, GM_ADDR rstd)
173- {173+ {
174- // GM BUFFER174+ // GM BUFFER
175- int64_t xOffset = blockIdx * blockFactor * numR;175+ int64_t xOffset = blockIdx * blockFactor * numR;
176- int64_t xLen = curBlockFactor * numR;176+ int64_t xLen = curBlockFactor * numR;
177- int64_t yOffset = blockIdx * blockFactor * numR;177+ int64_t yOffset = blockIdx * blockFactor * numR;
178- int64_t yLen = curBlockFactor * numR;178+ int64_t yLen = curBlockFactor * numR;
179- if constexpr (IsSameType<T_Y, int4b_t>::value) {179+ if constexpr (IsSameType<T_Y, int4b_t>::value) {
180- yOffset = yOffset / 2;180+ yOffset = yOffset / 2;
181- yLen = yLen / 2;181+ yLen = yLen / 2;
182- }182+ }
183- xGm.SetGlobalBuffer((__gm__ T_X*)x + xOffset, xLen);183+ xGm.SetGlobalBuffer((__gm__ T_X*)x + xOffset, xLen);
184- y1Gm.SetGlobalBuffer((__gm__ yDtype*)y1 + yOffset, yLen);184+ y1Gm.SetGlobalBuffer((__gm__ yDtype*)y1 + yOffset, yLen);
185- 185+ 
186- gammaGm.SetGlobalBuffer((__gm__ T_X*)gamma, numR);186+ gammaGm.SetGlobalBuffer((__gm__ T_X*)gamma, numR);
187- scales1Gm.SetGlobalBuffer((__gm__ T_SCALES*)scales1, numR);187+ scales1Gm.SetGlobalBuffer((__gm__ T_SCALES*)scales1, numR);
188- 188+ 
189- // gamma + scales1189+ // gamma + scales1
190- int64_t preloadDataSize = xGammaBetaAlign * sizeof(T_X) + scalesAlign * sizeof(T_SCALES);190+ int64_t preloadDataSize = xGammaBetaAlign * sizeof(T_X) + scalesAlign * sizeof(T_SCALES);
191- 191+ 
192- if (hasScales2) {192+ if (hasScales2) {
193- scales2Gm.SetGlobalBuffer((__gm__ T_SCALES*)scales2, numR);193+ scales2Gm.SetGlobalBuffer((__gm__ T_SCALES*)scales2, numR);
194- preloadDataSize = preloadDataSize + scalesAlign * sizeof(T_SCALES);194+ preloadDataSize = preloadDataSize + scalesAlign * sizeof(T_SCALES);
195- }195+ }
196- if (hasZeroPoints1) {196+ if (hasZeroPoints1) {
197- zeroPoints1Gm.SetGlobalBuffer((__gm__ T_ZEROPOINTS*)zeroPoints1, numR);197+ zeroPoints1Gm.SetGlobalBuffer((__gm__ T_ZEROPOINTS*)zeroPoints1, numR);
198- preloadDataSize = preloadDataSize + zeroPointsAlign * sizeof(T_ZEROPOINTS);198+ preloadDataSize = preloadDataSize + zeroPointsAlign * sizeof(T_ZEROPOINTS);
199- }199+ }
200- if (hasZeroPoints2) {200+ if (hasZeroPoints2) {
201- zeroPoints2Gm.SetGlobalBuffer((__gm__ T_ZEROPOINTS*)zeroPoints2, numR);201+ zeroPoints2Gm.SetGlobalBuffer((__gm__ T_ZEROPOINTS*)zeroPoints2, numR);
202- preloadDataSize = preloadDataSize + zeroPointsAlign * sizeof(T_ZEROPOINTS);202+ preloadDataSize = preloadDataSize + zeroPointsAlign * sizeof(T_ZEROPOINTS);
203- }203+ }
204- if (hasBeta) {204+ if (hasBeta) {
205- betaGm.SetGlobalBuffer((__gm__ T_X*)beta, numR);205+ betaGm.SetGlobalBuffer((__gm__ T_X*)beta, numR);
206- preloadDataSize = preloadDataSize + xGammaBetaAlign * sizeof(T_X);206+ preloadDataSize = preloadDataSize + xGammaBetaAlign * sizeof(T_X);
207- }207+ }
208- if (hasY2) {208+ if (hasY2) {
209- y2Gm.SetGlobalBuffer((__gm__ yDtype*)y2 + yOffset, yLen);209+ y2Gm.SetGlobalBuffer((__gm__ yDtype*)y2 + yOffset, yLen);
210- }210+ }
211- if (rstdFlag_ != 0) {211+ if (rstdFlag_ != 0) {
212- int64_t rstdOffset = blockIdx * blockFactor;212+ int64_t rstdOffset = blockIdx * blockFactor;
213- int64_t rstdLen = curBlockFactor;213+ int64_t rstdLen = curBlockFactor;
214- rstdGm.SetGlobalBuffer((__gm__ float*)rstd + rstdOffset, rstdLen);214+ rstdGm.SetGlobalBuffer((__gm__ float*)rstd + rstdOffset, rstdLen);
215- }215+ }
216- 216+ 
217- pipe_->InitBuffer(inQueueX, DOUBLE_BUFFER_NUM, ubFactor * xGammaBetaAlign * sizeof(T_X));217+ pipe_->InitBuffer(inQueueX, DOUBLE_BUFFER_NUM, ubFactor * xGammaBetaAlign * sizeof(T_X));
218- // preload data218+ // preload data
219- pipe_->InitBuffer(inQueueOhter, 1, preloadDataSize);219+ pipe_->InitBuffer(inQueueOhter, 1, preloadDataSize);
220- int64_t yQueueSize = ubFactor * yAlign * sizeof(yDtype);220+ int64_t yQueueSize = ubFactor * yAlign * sizeof(yDtype);
221- if constexpr (IsSameType<T_Y, int4b_t>::value) {221+ if constexpr (IsSameType<T_Y, int4b_t>::value) {
222- yQueueSize = yQueueSize / 2;222+ yQueueSize = yQueueSize / 2;
223- }223+ }
224- pipe_->InitBuffer(outQueueY1, DOUBLE_BUFFER_NUM, yQueueSize);224+ pipe_->InitBuffer(outQueueY1, DOUBLE_BUFFER_NUM, yQueueSize);
225- if (hasScales2) {225+ if (hasScales2) {
226- pipe_->InitBuffer(outQueueY2, DOUBLE_BUFFER_NUM, yQueueSize);226+ pipe_->InitBuffer(outQueueY2, DOUBLE_BUFFER_NUM, yQueueSize);
227- }227+ }
228- if (rstdFlag_ != 0) {228+ if (rstdFlag_ != 0) {
229- pipe_->InitBuffer(outQueueRstd, DOUBLE_BUFFER_NUM, rstdAlign * sizeof(float));229+ pipe_->InitBuffer(outQueueRstd, DOUBLE_BUFFER_NUM, rstdAlign * sizeof(float));
230- } else {230+ } else {
231- pipe_->InitBuffer(rstdBuf, rstdAlign * sizeof(float));231+ pipe_->InitBuffer(rstdBuf, rstdAlign * sizeof(float));
232- }232+ }
233- // reduceTmpBuffer233+ // reduceTmpBuffer
234- int64_t reduceTmpBufferSize = ubFactor *234+ int64_t reduceTmpBufferSize = ubFactor *
235- CeilDiv(CeilDiv(binaryAdd, static_cast<int64_t>(vectorLenB32)),235+ CeilDiv(CeilDiv(binaryAdd, static_cast<int64_t>(vectorLenB32)),
236- static_cast<int64_t>(blockSizeB32)) *236+ static_cast<int64_t>(blockSizeB32)) *
237- blockSizeB32;237+ blockSizeB32;
238- pipe_->InitBuffer(reduceTmpBuf, reduceTmpBufferSize);238+ pipe_->InitBuffer(reduceTmpBuf, reduceTmpBufferSize);
239- }239+ }
240- 240+ 
241- __aicore__ inline void Process()241+ __aicore__ inline void Process()
242- {242+ {
243- // copy other input (gamma scales zeropints beta)243+ // copy other input (gamma scales zeropints beta)
244- LocalTensor<uint8_t> otherLocal = inQueueOhter.AllocTensor<uint8_t>();244+ LocalTensor<uint8_t> otherLocal = inQueueOhter.AllocTensor<uint8_t>();
245- CopyInOhters(otherLocal);245+ CopyInOhters(otherLocal);
246- inQueueOhter.EnQue(otherLocal);246+ inQueueOhter.EnQue(otherLocal);
247- inQueueOhter.DeQue<uint8_t>();247+ inQueueOhter.DeQue<uint8_t>();
248- 248+ 
249- for (int64_t i = 0; i < curUbLoops; i++) {249+ for (int64_t i = 0; i < curUbLoops; i++) {
250- int64_t curUbFactor = (i == (curUbLoops - 1)) ? ubFactorTail : ubFactor; // ubFactorTail 尾部250+ int64_t curUbFactor = (i == (curUbLoops - 1)) ? ubFactorTail : ubFactor; // ubFactorTail 尾部
251- int64_t offsetBase = i * numR * ubFactor;251+ int64_t offsetBase = i * numR * ubFactor;
252- // x252+ // x
253- DataCopyPadExtParams<T_X> dataCopyPadExtParamsX;253+ DataCopyPadExtParams<T_X> dataCopyPadExtParamsX;
254- dataCopyPadExtParamsX.isPad = false;254+ dataCopyPadExtParamsX.isPad = false;
255- dataCopyPadExtParamsX.leftPadding = 0;255+ dataCopyPadExtParamsX.leftPadding = 0;
256- dataCopyPadExtParamsX.rightPadding = 0;256+ dataCopyPadExtParamsX.rightPadding = 0;
257- dataCopyPadExtParamsX.paddingValue = 0;257+ dataCopyPadExtParamsX.paddingValue = 0;
258- DataCopyExtParams copyInParamsX;258+ DataCopyExtParams copyInParamsX;
259- copyInParamsX.blockCount = curUbFactor;259+ copyInParamsX.blockCount = curUbFactor;
260- copyInParamsX.blockLen = numR * sizeof(T_X);260+ copyInParamsX.blockLen = numR * sizeof(T_X);
261- copyInParamsX.srcStride = 0;261+ copyInParamsX.srcStride = 0;
262- copyInParamsX.dstStride = 0;262+ copyInParamsX.dstStride = 0;
263- LocalTensor<T_X> xLocal = inQueueX.AllocTensor<T_X>();263+ LocalTensor<T_X> xLocal = inQueueX.AllocTensor<T_X>();
264- DataCopyPad(xLocal, xGm[offsetBase], copyInParamsX, dataCopyPadExtParamsX);264+ DataCopyPad(xLocal, xGm[offsetBase], copyInParamsX, dataCopyPadExtParamsX);
265- inQueueX.EnQue(xLocal);265+ inQueueX.EnQue(xLocal);
266- inQueueX.DeQue<T_X>();266+ inQueueX.DeQue<T_X>();
267- // compute square reduceSum & rstd267+ // compute square reduceSum & rstd
268- LocalTensor<float> reduceTmpLocal = reduceTmpBuf.Get<float>();268+ LocalTensor<float> reduceTmpLocal = reduceTmpBuf.Get<float>();
269- LocalTensor<float> rstdLocal;269+ LocalTensor<float> rstdLocal;
270- if (rstdFlag_ != 0) {270+ if (rstdFlag_ != 0) {
271- rstdLocal = outQueueRstd.AllocTensor<float>();271+ rstdLocal = outQueueRstd.AllocTensor<float>();
272- } else {272+ } else {
273- rstdLocal = rstdBuf.Get<float>();273+ rstdLocal = rstdBuf.Get<float>();
274- }274+ }
275- NormCommon::NormCommonRegbase::CalculateSquareReduceSum<T_X>(275+ NormCommon::NormCommonRegbase::CalculateSquareReduceSum<T_X>(
276- xLocal, rstdLocal, reduceTmpLocal, static_cast<uint16_t>(curUbFactor),276+ xLocal, rstdLocal, reduceTmpLocal, static_cast<uint16_t>(curUbFactor),
277- static_cast<uint32_t>(xGammaBetaAlign), static_cast<uint32_t>(numR), static_cast<uint32_t>(binaryAdd),277+ static_cast<uint32_t>(xGammaBetaAlign), static_cast<uint32_t>(numR), static_cast<uint32_t>(binaryAdd),
278- static_cast<uint32_t>(blockSizeB32), static_cast<uint32_t>(xGammaBetaAlign));278+ static_cast<uint32_t>(blockSizeB32), static_cast<uint32_t>(xGammaBetaAlign));
279- NormCommon::ComputeRstdNewtonRaphson<false, true>(rstdLocal, rstdLocal, static_cast<uint32_t>(curUbFactor),279+ NormCommon::ComputeRstdNewtonRaphson<false, true>(rstdLocal, rstdLocal, static_cast<uint32_t>(curUbFactor),
280- epsilon, avgFactor, vectorLenB32);280+ epsilon, avgFactor, vectorLenB32);
281- if (rstdFlag_ != 0) {281+ if (rstdFlag_ != 0) {
282- outQueueRstd.EnQue(rstdLocal);282+ outQueueRstd.EnQue(rstdLocal);
283- rstdLocal = outQueueRstd.DeQue<float>();283+ rstdLocal = outQueueRstd.DeQue<float>();
284- DataCopyExtParams copyOutParamsRstd;284+ DataCopyExtParams copyOutParamsRstd;
285- copyOutParamsRstd.blockCount = 1;285+ copyOutParamsRstd.blockCount = 1;
286- copyOutParamsRstd.blockLen = curUbFactor * sizeof(float);286+ copyOutParamsRstd.blockLen = curUbFactor * sizeof(float);
287- copyOutParamsRstd.srcStride = 0;287+ copyOutParamsRstd.srcStride = 0;
288- copyOutParamsRstd.dstStride = 0;288+ copyOutParamsRstd.dstStride = 0;
289- DataCopyPad(rstdGm[i * ubFactor], rstdLocal, copyOutParamsRstd);289+ DataCopyPad(rstdGm[i * ubFactor], rstdLocal, copyOutParamsRstd);
290- }290+ }
291- 291+ 
292- LocalTensor<yDtype> y1Local = outQueueY1.AllocTensor<yDtype>();292+ LocalTensor<yDtype> y1Local = outQueueY1.AllocTensor<yDtype>();
293- LocalTensor<yDtype> y2Local;293+ LocalTensor<yDtype> y2Local;
294- if (hasY2) {294+ if (hasY2) {
295- y2Local = outQueueY2.AllocTensor<yDtype>();295+ y2Local = outQueueY2.AllocTensor<yDtype>();
296- }296+ }
297- 297+ 
298- // compute quant298+ // compute quant
299- SetOverflowMode<T_Y>(0);299+ SetOverflowMode<T_Y>(0);
300- QuantRoute(optionMask, xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local, scales2Local,300+ QuantRoute(optionMask, xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local, scales2Local,
301- zeroPoints1Local, zeroPoints2Local, y1Local, y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,301+ zeroPoints1Local, zeroPoints2Local, y1Local, y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
302- scalesAlign, zeroPointsAlign, yAlign);302+ scalesAlign, zeroPointsAlign, yAlign);
303- SetOverflowMode<T_Y>(oriOverflowMode);303+ SetOverflowMode<T_Y>(oriOverflowMode);
304- inQueueX.FreeTensor(xLocal);304+ inQueueX.FreeTensor(xLocal);
305- if (rstdFlag_ != 0) {305+ if (rstdFlag_ != 0) {
306- outQueueRstd.FreeTensor(rstdLocal);306+ outQueueRstd.FreeTensor(rstdLocal);
307- }307+ }
308- 308+ 
309- outQueueY1.EnQue(y1Local);309+ outQueueY1.EnQue(y1Local);
310- outQueueY1.DeQue<yDtype>();310+ outQueueY1.DeQue<yDtype>();
311- 311+ 
312- int64_t yOffsetBase = offsetBase;312+ int64_t yOffsetBase = offsetBase;
313- int64_t yBlockLen = numR * sizeof(yDtype);313+ int64_t yBlockLen = numR * sizeof(yDtype);
314- if constexpr (IsSameType<T_Y, int4b_t>::value) {314+ if constexpr (IsSameType<T_Y, int4b_t>::value) {
315- yOffsetBase = yOffsetBase / 2;315+ yOffsetBase = yOffsetBase / 2;
316- yBlockLen = yBlockLen / 2;316+ yBlockLen = yBlockLen / 2;
317- }317+ }
318- DataCopyExtParams copyOutParamsY1;318+ DataCopyExtParams copyOutParamsY1;
319- copyOutParamsY1.blockCount = curUbFactor;319+ copyOutParamsY1.blockCount = curUbFactor;
320- copyOutParamsY1.blockLen = yBlockLen;320+ copyOutParamsY1.blockLen = yBlockLen;
321- copyOutParamsY1.srcStride = 0;321+ copyOutParamsY1.srcStride = 0;
322- copyOutParamsY1.dstStride = 0;322+ copyOutParamsY1.dstStride = 0;
323- DataCopyPad(y1Gm[yOffsetBase], y1Local, copyOutParamsY1);323+ DataCopyPad(y1Gm[yOffsetBase], y1Local, copyOutParamsY1);
324- outQueueY1.FreeTensor(y1Local);324+ outQueueY1.FreeTensor(y1Local);
325- 325+ 
326- if (hasY2) {326+ if (hasY2) {
327- outQueueY2.EnQue(y2Local);327+ outQueueY2.EnQue(y2Local);
328- outQueueY2.DeQue<yDtype>();328+ outQueueY2.DeQue<yDtype>();
329- DataCopyExtParams copyOutParamsY2;329+ DataCopyExtParams copyOutParamsY2;
330- copyOutParamsY2.blockCount = curUbFactor;330+ copyOutParamsY2.blockCount = curUbFactor;
331- copyOutParamsY2.blockLen = yBlockLen;331+ copyOutParamsY2.blockLen = yBlockLen;
332- copyOutParamsY2.srcStride = 0;332+ copyOutParamsY2.srcStride = 0;
333- copyOutParamsY2.dstStride = 0;333+ copyOutParamsY2.dstStride = 0;
334- DataCopyPad(y2Gm[yOffsetBase], y2Local, copyOutParamsY2);334+ DataCopyPad(y2Gm[yOffsetBase], y2Local, copyOutParamsY2);
335- outQueueY2.FreeTensor(y2Local);335+ outQueueY2.FreeTensor(y2Local);
336- }336+ }
337- }337+ }
338- inQueueOhter.FreeTensor(otherLocal);338+ inQueueOhter.FreeTensor(otherLocal);
339- }339+ }
340- 340+ 
341-private:341+private:
342- __aicore__ inline void CopyInOhters(LocalTensor<uint8_t> otherLocal)342+ __aicore__ inline void CopyInOhters(LocalTensor<uint8_t> otherLocal)
343- {343+ {
344- uint32_t localOffset = 0;344+ uint32_t localOffset = 0;
345- // LocalTensor<T_X> gammaLocal345+ // LocalTensor<T_X> gammaLocal
346- gammaLocal = otherLocal[localOffset].ReinterpretCast<T_X>();346+ gammaLocal = otherLocal[localOffset].ReinterpretCast<T_X>();
347- localOffset = localOffset + xGammaBetaAlign * sizeof(T_X);347+ localOffset = localOffset + xGammaBetaAlign * sizeof(T_X);
348- DataCopyPadExtParams<T_X> dataCopyPadExtParamsGamma;348+ DataCopyPadExtParams<T_X> dataCopyPadExtParamsGamma;
349- dataCopyPadExtParamsGamma.isPad = false;349+ dataCopyPadExtParamsGamma.isPad = false;
350- dataCopyPadExtParamsGamma.leftPadding = 0;350+ dataCopyPadExtParamsGamma.leftPadding = 0;
351- dataCopyPadExtParamsGamma.rightPadding = 0;351+ dataCopyPadExtParamsGamma.rightPadding = 0;
352- dataCopyPadExtParamsGamma.paddingValue = 0;352+ dataCopyPadExtParamsGamma.paddingValue = 0;
353- DataCopyExtParams copyInParamsGamma;353+ DataCopyExtParams copyInParamsGamma;
354- copyInParamsGamma.blockCount = 1;354+ copyInParamsGamma.blockCount = 1;
355- copyInParamsGamma.blockLen = numR * sizeof(T_X);355+ copyInParamsGamma.blockLen = numR * sizeof(T_X);
356- copyInParamsGamma.srcStride = 0;356+ copyInParamsGamma.srcStride = 0;
357- copyInParamsGamma.dstStride = 0;357+ copyInParamsGamma.dstStride = 0;
358- DataCopyPad(gammaLocal, gammaGm, copyInParamsGamma, dataCopyPadExtParamsGamma);358+ DataCopyPad(gammaLocal, gammaGm, copyInParamsGamma, dataCopyPadExtParamsGamma);
359- 359+ 
360- // LocalTensor<T_SCALES> scales1Local;360+ // LocalTensor<T_SCALES> scales1Local;
361- scales1Local = otherLocal[localOffset].ReinterpretCast<T_SCALES>();361+ scales1Local = otherLocal[localOffset].ReinterpretCast<T_SCALES>();
362- localOffset = localOffset + scalesAlign * sizeof(T_SCALES);362+ localOffset = localOffset + scalesAlign * sizeof(T_SCALES);
363- DataCopyPadExtParams<T_SCALES> dataCopyPadExtParamsScales;363+ DataCopyPadExtParams<T_SCALES> dataCopyPadExtParamsScales;
364- dataCopyPadExtParamsScales.isPad = false;364+ dataCopyPadExtParamsScales.isPad = false;
365- dataCopyPadExtParamsScales.leftPadding = 0;365+ dataCopyPadExtParamsScales.leftPadding = 0;
366- dataCopyPadExtParamsScales.rightPadding = 0;366+ dataCopyPadExtParamsScales.rightPadding = 0;
367- dataCopyPadExtParamsScales.paddingValue = 0;367+ dataCopyPadExtParamsScales.paddingValue = 0;
368- DataCopyExtParams copyInParamsScales;368+ DataCopyExtParams copyInParamsScales;
369- copyInParamsScales.blockCount = 1;369+ copyInParamsScales.blockCount = 1;
370- if (numQ == 1) {370+ if (numQ == 1) {
371- copyInParamsScales.blockLen = sizeof(T_SCALES);371+ copyInParamsScales.blockLen = sizeof(T_SCALES);
372- } else {372+ } else {
373- copyInParamsScales.blockLen = numR * sizeof(T_SCALES);373+ copyInParamsScales.blockLen = numR * sizeof(T_SCALES);
374- }374+ }
375- copyInParamsScales.srcStride = 0;375+ copyInParamsScales.srcStride = 0;
376- copyInParamsScales.dstStride = 0;376+ copyInParamsScales.dstStride = 0;
377- DataCopyPad(scales1Local, scales1Gm, copyInParamsScales, dataCopyPadExtParamsScales);377+ DataCopyPad(scales1Local, scales1Gm, copyInParamsScales, dataCopyPadExtParamsScales);
378- 378+ 
379- // zeroPoints 代码379+ // zeroPoints 代码
380- DataCopyPadExtParams<T_ZEROPOINTS> dataCopyPadExtParamszeroPoints;380+ DataCopyPadExtParams<T_ZEROPOINTS> dataCopyPadExtParamszeroPoints;
381- dataCopyPadExtParamszeroPoints.isPad = false;381+ dataCopyPadExtParamszeroPoints.isPad = false;
382- dataCopyPadExtParamszeroPoints.leftPadding = 0;382+ dataCopyPadExtParamszeroPoints.leftPadding = 0;
383- dataCopyPadExtParamszeroPoints.rightPadding = 0;383+ dataCopyPadExtParamszeroPoints.rightPadding = 0;
384- dataCopyPadExtParamszeroPoints.paddingValue = 0;384+ dataCopyPadExtParamszeroPoints.paddingValue = 0;
385- DataCopyExtParams copyInParamszeroPoints;385+ DataCopyExtParams copyInParamszeroPoints;
386- copyInParamszeroPoints.blockCount = 1;386+ copyInParamszeroPoints.blockCount = 1;
387- if (numQ == 1) {387+ if (numQ == 1) {
388- copyInParamszeroPoints.blockLen = sizeof(T_ZEROPOINTS);388+ copyInParamszeroPoints.blockLen = sizeof(T_ZEROPOINTS);
389- } else {389+ } else {
390- copyInParamszeroPoints.blockLen = numR * sizeof(T_ZEROPOINTS);390+ copyInParamszeroPoints.blockLen = numR * sizeof(T_ZEROPOINTS);
391- }391+ }
392- copyInParamszeroPoints.srcStride = 0;392+ copyInParamszeroPoints.srcStride = 0;
393- copyInParamszeroPoints.dstStride = 0;393+ copyInParamszeroPoints.dstStride = 0;
394- 394+ 
395- if (hasScales2) {395+ if (hasScales2) {
396- // LocalTensor<T_SCALES> scales2Local;396+ // LocalTensor<T_SCALES> scales2Local;
397- scales2Local = otherLocal[localOffset].ReinterpretCast<T_SCALES>();397+ scales2Local = otherLocal[localOffset].ReinterpretCast<T_SCALES>();
398- localOffset = localOffset + scalesAlign * sizeof(T_SCALES);398+ localOffset = localOffset + scalesAlign * sizeof(T_SCALES);
399- DataCopyPad(scales2Local, scales2Gm, copyInParamsScales, dataCopyPadExtParamsScales);399+ DataCopyPad(scales2Local, scales2Gm, copyInParamsScales, dataCopyPadExtParamsScales);
400- }400+ }
401- 401+ 
402- if (hasZeroPoints1) {402+ if (hasZeroPoints1) {
403- // LocalTensor<T_ZEROPOINTS> zeroPoints1Local;403+ // LocalTensor<T_ZEROPOINTS> zeroPoints1Local;
404- zeroPoints1Local = otherLocal[localOffset].ReinterpretCast<T_ZEROPOINTS>();404+ zeroPoints1Local = otherLocal[localOffset].ReinterpretCast<T_ZEROPOINTS>();
405- localOffset = localOffset + zeroPointsAlign * sizeof(T_ZEROPOINTS);405+ localOffset = localOffset + zeroPointsAlign * sizeof(T_ZEROPOINTS);
406- DataCopyPad(zeroPoints1Local, zeroPoints1Gm, copyInParamszeroPoints, dataCopyPadExtParamszeroPoints);406+ DataCopyPad(zeroPoints1Local, zeroPoints1Gm, copyInParamszeroPoints, dataCopyPadExtParamszeroPoints);
407- }407+ }
408- 408+ 
409- if (hasZeroPoints2) {409+ if (hasZeroPoints2) {
410- // LocalTensor<T_ZEROPOINTS> zeroPoints2Local;410+ // LocalTensor<T_ZEROPOINTS> zeroPoints2Local;
411- zeroPoints2Local = otherLocal[localOffset].ReinterpretCast<T_ZEROPOINTS>();411+ zeroPoints2Local = otherLocal[localOffset].ReinterpretCast<T_ZEROPOINTS>();
412- localOffset = localOffset + zeroPointsAlign * sizeof(T_ZEROPOINTS);412+ localOffset = localOffset + zeroPointsAlign * sizeof(T_ZEROPOINTS);
413- DataCopyPad(zeroPoints2Local, zeroPoints2Gm, copyInParamszeroPoints, dataCopyPadExtParamszeroPoints);413+ DataCopyPad(zeroPoints2Local, zeroPoints2Gm, copyInParamszeroPoints, dataCopyPadExtParamszeroPoints);
414- }414+ }
415- 415+ 
416- if (hasBeta) {416+ if (hasBeta) {
417- // LocalTensor<T_X> betaLocal;417+ // LocalTensor<T_X> betaLocal;
418- betaLocal = otherLocal[localOffset].ReinterpretCast<T_X>();418+ betaLocal = otherLocal[localOffset].ReinterpretCast<T_X>();
419- DataCopyPad(betaLocal, betaGm, copyInParamsGamma, dataCopyPadExtParamsGamma);419+ DataCopyPad(betaLocal, betaGm, copyInParamsGamma, dataCopyPadExtParamsGamma);
420- }420+ }
421- }421+ }
422- 422+ 
423- __aicore__ inline void QuantRoute(uint32_t optionMask, LocalTensor<T_X> xLocal, LocalTensor<float> rstdLocal,423+ __aicore__ inline void QuantRoute(uint32_t optionMask, LocalTensor<T_X> xLocal, LocalTensor<float> rstdLocal,
424- LocalTensor<T_X> gammaLocal, LocalTensor<T_X> betaLocal,424+ LocalTensor<T_X> gammaLocal, LocalTensor<T_X> betaLocal,
425- LocalTensor<T_SCALES> scales1Local, LocalTensor<T_SCALES> scales2Local,425+ LocalTensor<T_SCALES> scales1Local, LocalTensor<T_SCALES> scales2Local,
426- LocalTensor<T_ZEROPOINTS> zeroPoints1Local,426+ LocalTensor<T_ZEROPOINTS> zeroPoints1Local,
427- LocalTensor<T_ZEROPOINTS> zeroPoints2Local, LocalTensor<yDtype> y1Local,427+ LocalTensor<T_ZEROPOINTS> zeroPoints2Local, LocalTensor<yDtype> y1Local,
428- LocalTensor<yDtype> y2Local, int64_t curUbFactor, int64_t numR, int64_t numQ,428+ LocalTensor<yDtype> y2Local, int64_t curUbFactor, int64_t numR, int64_t numQ,
429- int64_t xGammaBetaAlign, int64_t scalesAlign, int64_t zeroPointsAlign,429+ int64_t xGammaBetaAlign, int64_t scalesAlign, int64_t zeroPointsAlign,
430- int64_t yAlign)430+ int64_t yAlign)
431- {431+ {
432- // compute quant432+ // compute quant
433- if (!isScaleDiv) {433+ if (!isScaleDiv) {
434- if (optionMask == 0b1111) {434+ if (optionMask == 0b1111) {
435- ComputeQuant<true, true, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,435+ ComputeQuant<true, true, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
436- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,436+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
437- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,437+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
438- scalesAlign, zeroPointsAlign, yAlign);438+ scalesAlign, zeroPointsAlign, yAlign);
439- } else if (optionMask == 0b0111) {439+ } else if (optionMask == 0b0111) {
440- ComputeQuant<false, true, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,440+ ComputeQuant<false, true, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
441- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,441+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
442- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,442+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
443- scalesAlign, zeroPointsAlign, yAlign);443+ scalesAlign, zeroPointsAlign, yAlign);
444- } else if (optionMask == 0b1011) {444+ } else if (optionMask == 0b1011) {
445- ComputeQuant<true, false, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,445+ ComputeQuant<true, false, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
446- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,446+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
447- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,447+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
448- scalesAlign, zeroPointsAlign, yAlign);448+ scalesAlign, zeroPointsAlign, yAlign);
449- } else if (optionMask == 0b0011) {449+ } else if (optionMask == 0b0011) {
450- ComputeQuant<false, false, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,450+ ComputeQuant<false, false, true, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
451- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,451+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
452- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,452+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
453- scalesAlign, zeroPointsAlign, yAlign);453+ scalesAlign, zeroPointsAlign, yAlign);
454- } else if (optionMask == 0b1101) {454+ } else if (optionMask == 0b1101) {
455- ComputeQuant<true, true, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,455+ ComputeQuant<true, true, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
456- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,456+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
457- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,457+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
458- scalesAlign, zeroPointsAlign, yAlign);458+ scalesAlign, zeroPointsAlign, yAlign);
459- } else if (optionMask == 0b0101) {459+ } else if (optionMask == 0b0101) {
460- ComputeQuant<false, true, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,460+ ComputeQuant<false, true, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
461- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,461+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
462- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,462+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
463- scalesAlign, zeroPointsAlign, yAlign);463+ scalesAlign, zeroPointsAlign, yAlign);
464- } else if (optionMask == 0b1001) {464+ } else if (optionMask == 0b1001) {
465- ComputeQuant<true, false, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,465+ ComputeQuant<true, false, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
466- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,466+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
467- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,467+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
468- scalesAlign, zeroPointsAlign, yAlign);468+ scalesAlign, zeroPointsAlign, yAlign);
469- } else if (optionMask == 0b0001) {469+ } else if (optionMask == 0b0001) {
470- ComputeQuant<false, false, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,470+ ComputeQuant<false, false, false, true, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
471- scales2Local, zeroPoints1Local, zeroPoints2Local,471+ scales2Local, zeroPoints1Local, zeroPoints2Local,
472- y1Local, y2Local, curUbFactor, numR, numQ,472+ y1Local, y2Local, curUbFactor, numR, numQ,
473- xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);473+ xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);
474- } else if (optionMask == 0b1110) {474+ } else if (optionMask == 0b1110) {
475- ComputeQuant<true, true, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,475+ ComputeQuant<true, true, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
476- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,476+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
477- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,477+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
478- scalesAlign, zeroPointsAlign, yAlign);478+ scalesAlign, zeroPointsAlign, yAlign);
479- } else if (optionMask == 0b0110) {479+ } else if (optionMask == 0b0110) {
480- ComputeQuant<false, true, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,480+ ComputeQuant<false, true, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
481- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,481+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
482- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,482+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
483- scalesAlign, zeroPointsAlign, yAlign);483+ scalesAlign, zeroPointsAlign, yAlign);
484- } else if (optionMask == 0b1010) {484+ } else if (optionMask == 0b1010) {
485- ComputeQuant<true, false, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,485+ ComputeQuant<true, false, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
486- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,486+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
487- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,487+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
488- scalesAlign, zeroPointsAlign, yAlign);488+ scalesAlign, zeroPointsAlign, yAlign);
489- } else if (optionMask == 0b0010) {489+ } else if (optionMask == 0b0010) {
490- ComputeQuant<false, false, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,490+ ComputeQuant<false, false, true, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
491- scales2Local, zeroPoints1Local, zeroPoints2Local,491+ scales2Local, zeroPoints1Local, zeroPoints2Local,
492- y1Local, y2Local, curUbFactor, numR, numQ,492+ y1Local, y2Local, curUbFactor, numR, numQ,
493- xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);493+ xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);
494- } else if (optionMask == 0b1100) {494+ } else if (optionMask == 0b1100) {
495- ComputeQuant<true, true, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,495+ ComputeQuant<true, true, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
496- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,496+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
497- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,497+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
498- scalesAlign, zeroPointsAlign, yAlign);498+ scalesAlign, zeroPointsAlign, yAlign);
499- } else if (optionMask == 0b0100) {499+ } else if (optionMask == 0b0100) {
500- ComputeQuant<false, true, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,500+ ComputeQuant<false, true, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
501- scales2Local, zeroPoints1Local, zeroPoints2Local,501+ scales2Local, zeroPoints1Local, zeroPoints2Local,
502- y1Local, y2Local, curUbFactor, numR, numQ,502+ y1Local, y2Local, curUbFactor, numR, numQ,
503- xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);503+ xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);
504- } else if (optionMask == 0b1000) {504+ } else if (optionMask == 0b1000) {
505- ComputeQuant<true, false, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,505+ ComputeQuant<true, false, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
506- scales2Local, zeroPoints1Local, zeroPoints2Local,506+ scales2Local, zeroPoints1Local, zeroPoints2Local,
507- y1Local, y2Local, curUbFactor, numR, numQ,507+ y1Local, y2Local, curUbFactor, numR, numQ,
508- xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);508+ xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);
509- } else if (optionMask == 0b0000) {509+ } else if (optionMask == 0b0000) {
510- ComputeQuant<false, false, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,510+ ComputeQuant<false, false, false, false, false>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
511- scales2Local, zeroPoints1Local, zeroPoints2Local,511+ scales2Local, zeroPoints1Local, zeroPoints2Local,
512- y1Local, y2Local, curUbFactor, numR, numQ,512+ y1Local, y2Local, curUbFactor, numR, numQ,
513- xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);513+ xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);
514- }514+ }
515- } else {515+ } else {
516- if (optionMask == 0b1111) {516+ if (optionMask == 0b1111) {
517- ComputeQuant<true, true, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,517+ ComputeQuant<true, true, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
518- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,518+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
519- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,519+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
520- scalesAlign, zeroPointsAlign, yAlign);520+ scalesAlign, zeroPointsAlign, yAlign);
521- } else if (optionMask == 0b0111) {521+ } else if (optionMask == 0b0111) {
522- ComputeQuant<false, true, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,522+ ComputeQuant<false, true, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
523- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,523+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
524- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,524+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
525- scalesAlign, zeroPointsAlign, yAlign);525+ scalesAlign, zeroPointsAlign, yAlign);
526- } else if (optionMask == 0b1011) {526+ } else if (optionMask == 0b1011) {
527- ComputeQuant<true, false, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,527+ ComputeQuant<true, false, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
528- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,528+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
529- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,529+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
530- scalesAlign, zeroPointsAlign, yAlign);530+ scalesAlign, zeroPointsAlign, yAlign);
531- } else if (optionMask == 0b0011) {531+ } else if (optionMask == 0b0011) {
532- ComputeQuant<false, false, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,532+ ComputeQuant<false, false, true, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
533- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,533+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
534- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,534+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
535- scalesAlign, zeroPointsAlign, yAlign);535+ scalesAlign, zeroPointsAlign, yAlign);
536- } else if (optionMask == 0b1101) {536+ } else if (optionMask == 0b1101) {
537- ComputeQuant<true, true, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,537+ ComputeQuant<true, true, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
538- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,538+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
539- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,539+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
540- scalesAlign, zeroPointsAlign, yAlign);540+ scalesAlign, zeroPointsAlign, yAlign);
541- } else if (optionMask == 0b0101) {541+ } else if (optionMask == 0b0101) {
542- ComputeQuant<false, true, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,542+ ComputeQuant<false, true, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
543- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,543+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
544- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,544+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
545- scalesAlign, zeroPointsAlign, yAlign);545+ scalesAlign, zeroPointsAlign, yAlign);
546- } else if (optionMask == 0b1001) {546+ } else if (optionMask == 0b1001) {
547- ComputeQuant<true, false, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,547+ ComputeQuant<true, false, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
548- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,548+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
549- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,549+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
550- scalesAlign, zeroPointsAlign, yAlign);550+ scalesAlign, zeroPointsAlign, yAlign);
551- } else if (optionMask == 0b0001) {551+ } else if (optionMask == 0b0001) {
552- ComputeQuant<false, false, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,552+ ComputeQuant<false, false, false, true, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
553- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,553+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
554- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,554+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
555- scalesAlign, zeroPointsAlign, yAlign);555+ scalesAlign, zeroPointsAlign, yAlign);
556- } else if (optionMask == 0b1110) {556+ } else if (optionMask == 0b1110) {
557- ComputeQuant<true, true, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,557+ ComputeQuant<true, true, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
558- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,558+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
559- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,559+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
560- scalesAlign, zeroPointsAlign, yAlign);560+ scalesAlign, zeroPointsAlign, yAlign);
561- } else if (optionMask == 0b0110) {561+ } else if (optionMask == 0b0110) {
562- ComputeQuant<false, true, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,562+ ComputeQuant<false, true, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
563- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,563+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
564- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,564+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
565- scalesAlign, zeroPointsAlign, yAlign);565+ scalesAlign, zeroPointsAlign, yAlign);
566- } else if (optionMask == 0b1010) {566+ } else if (optionMask == 0b1010) {
567- ComputeQuant<true, false, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,567+ ComputeQuant<true, false, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
568- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,568+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
569- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,569+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
570- scalesAlign, zeroPointsAlign, yAlign);570+ scalesAlign, zeroPointsAlign, yAlign);
571- } else if (optionMask == 0b0010) {571+ } else if (optionMask == 0b0010) {
572- ComputeQuant<false, false, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,572+ ComputeQuant<false, false, true, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
573- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,573+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
574- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,574+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
575- scalesAlign, zeroPointsAlign, yAlign);575+ scalesAlign, zeroPointsAlign, yAlign);
576- } else if (optionMask == 0b1100) {576+ } else if (optionMask == 0b1100) {
577- ComputeQuant<true, true, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,577+ ComputeQuant<true, true, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
578- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,578+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
579- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,579+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
580- scalesAlign, zeroPointsAlign, yAlign);580+ scalesAlign, zeroPointsAlign, yAlign);
581- } else if (optionMask == 0b0100) {581+ } else if (optionMask == 0b0100) {
582- ComputeQuant<false, true, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,582+ ComputeQuant<false, true, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
583- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,583+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
584- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,584+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
585- scalesAlign, zeroPointsAlign, yAlign);585+ scalesAlign, zeroPointsAlign, yAlign);
586- } else if (optionMask == 0b1000) {586+ } else if (optionMask == 0b1000) {
587- ComputeQuant<true, false, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,587+ ComputeQuant<true, false, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
588- scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,588+ scales2Local, zeroPoints1Local, zeroPoints2Local, y1Local,
589- y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,589+ y2Local, curUbFactor, numR, numQ, xGammaBetaAlign,
590- scalesAlign, zeroPointsAlign, yAlign);590+ scalesAlign, zeroPointsAlign, yAlign);
591- } else if (optionMask == 0b0000) {591+ } else if (optionMask == 0b0000) {
592- ComputeQuant<false, false, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,592+ ComputeQuant<false, false, false, false, true>(xLocal, rstdLocal, gammaLocal, betaLocal, scales1Local,
593- scales2Local, zeroPoints1Local, zeroPoints2Local,593+ scales2Local, zeroPoints1Local, zeroPoints2Local,
594- y1Local, y2Local, curUbFactor, numR, numQ,594+ y1Local, y2Local, curUbFactor, numR, numQ,
595- xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);595+ xGammaBetaAlign, scalesAlign, zeroPointsAlign, yAlign);
596- }596+ }
597- }597+ }
598- }598+ }
599- 599+ 
600- template <bool HAS_BETA, bool HAS_ZEROPINTS2, bool HAS_ZEROPINTS1, bool HAS_SCALES2, bool IS_SCALES_DIV>600+ template <bool HAS_BETA, bool HAS_ZEROPINTS2, bool HAS_ZEROPINTS1, bool HAS_SCALES2, bool IS_SCALES_DIV>
601- __aicore__ inline void ComputeQuant(LocalTensor<T_X> xLocal, LocalTensor<float> rstdLocal,601+ __aicore__ inline void ComputeQuant(LocalTensor<T_X> xLocal, LocalTensor<float> rstdLocal,
602- LocalTensor<T_X> gammaLocal, LocalTensor<T_X> betaLocal,602+ LocalTensor<T_X> gammaLocal, LocalTensor<T_X> betaLocal,
603- LocalTensor<T_SCALES> scales1Local, LocalTensor<T_SCALES> scales2Local,603+ LocalTensor<T_SCALES> scales1Local, LocalTensor<T_SCALES> scales2Local,
604- LocalTensor<T_ZEROPOINTS> zeroPoints1Local,604+ LocalTensor<T_ZEROPOINTS> zeroPoints1Local,
605- LocalTensor<T_ZEROPOINTS> zeroPoints2Local, LocalTensor<yDtype> y1Local,605+ LocalTensor<T_ZEROPOINTS> zeroPoints2Local, LocalTensor<yDtype> y1Local,
606- LocalTensor<yDtype> y2Local, int64_t curUbFactor, int64_t numR, int64_t numQ,606+ LocalTensor<yDtype> y2Local, int64_t curUbFactor, int64_t numR, int64_t numQ,
607- int64_t xGammaBetaAlign, int64_t scalesAlign, int64_t zeroPointsAlign,607+ int64_t xGammaBetaAlign, int64_t scalesAlign, int64_t zeroPointsAlign,
608- int64_t yAlign)608+ int64_t yAlign)
609- {609+ {
610- uint16_t loopsA = static_cast<uint16_t>(curUbFactor);610+ uint16_t loopsA = static_cast<uint16_t>(curUbFactor);
611- uint16_t loopsR = static_cast<uint16_t>(CeilDiv(static_cast<uint32_t>(numR), vectorLenB32));611+ uint16_t loopsR = static_cast<uint16_t>(CeilDiv(static_cast<uint32_t>(numR), vectorLenB32));
612- uint32_t sregR = static_cast<uint16_t>(numR);612+ uint32_t sregR = static_cast<uint16_t>(numR);
613- uint32_t sregxGammaBetaAlign = static_cast<uint16_t>(xGammaBetaAlign);613+ uint32_t sregxGammaBetaAlign = static_cast<uint16_t>(xGammaBetaAlign);
614- uint32_t sregyAlign = static_cast<uint16_t>(yAlign);614+ uint32_t sregyAlign = static_cast<uint16_t>(yAlign);
615- __local_mem__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();615+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
616- __local_mem__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();616+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
617- __local_mem__ T_X* gammaAddr = (__ubuf__ T_X*)gammaLocal.GetPhyAddr();617+ __ubuf__ T_X* gammaAddr = (__ubuf__ T_X*)gammaLocal.GetPhyAddr();
618- __local_mem__ T_SCALES* scales1Addr = (__ubuf__ T_SCALES*)scales1Local.GetPhyAddr();618+ __ubuf__ T_SCALES* scales1Addr = (__ubuf__ T_SCALES*)scales1Local.GetPhyAddr();
619- __local_mem__ T_ZEROPOINTS* zeroPoints1Addr;619+ __ubuf__ T_ZEROPOINTS* zeroPoints1Addr;
620- __local_mem__ T_SCALES* scales2Addr;620+ __ubuf__ T_SCALES* scales2Addr;
621- __local_mem__ T_ZEROPOINTS* zeroPoints2Addr;621+ __ubuf__ T_ZEROPOINTS* zeroPoints2Addr;
622- __local_mem__ T_X* betaAddr;622+ __ubuf__ T_X* betaAddr;
623- __local_mem__ yDtype* y1Addr;623+ __ubuf__ yDtype* y1Addr;
624- __local_mem__ yDtype* y2Addr;624+ __ubuf__ yDtype* y2Addr;
625- 625+ 
626- if constexpr (HAS_ZEROPINTS1) {626+ if constexpr (HAS_ZEROPINTS1) {
627- zeroPoints1Addr = (__ubuf__ T_ZEROPOINTS*)zeroPoints1Local.GetPhyAddr();627+ zeroPoints1Addr = (__ubuf__ T_ZEROPOINTS*)zeroPoints1Local.GetPhyAddr();
628- }628+ }
629- if constexpr (HAS_SCALES2) {629+ if constexpr (HAS_SCALES2) {
630- scales2Addr = (__ubuf__ T_SCALES*)scales2Local.GetPhyAddr();630+ scales2Addr = (__ubuf__ T_SCALES*)scales2Local.GetPhyAddr();
631- }631+ }
632- if constexpr (HAS_ZEROPINTS2) {632+ if constexpr (HAS_ZEROPINTS2) {
633- zeroPoints2Addr = (__ubuf__ T_ZEROPOINTS*)zeroPoints2Local.GetPhyAddr();633+ zeroPoints2Addr = (__ubuf__ T_ZEROPOINTS*)zeroPoints2Local.GetPhyAddr();
634- }634+ }
635- if constexpr (HAS_BETA) {635+ if constexpr (HAS_BETA) {
636- betaAddr = (__ubuf__ T_X*)betaLocal.GetPhyAddr();636+ betaAddr = (__ubuf__ T_X*)betaLocal.GetPhyAddr();
637- }637+ }
638- 638+ 
639- y1Addr = (__ubuf__ yDtype*)y1Local.GetPhyAddr();639+ y1Addr = (__ubuf__ yDtype*)y1Local.GetPhyAddr();
640- if constexpr ((HAS_ZEROPINTS2 || HAS_SCALES2)) {640+ if constexpr ((HAS_ZEROPINTS2 || HAS_SCALES2)) {
641- y2Addr = (__ubuf__ yDtype*)y2Local.GetPhyAddr();641+ y2Addr = (__ubuf__ yDtype*)y2Local.GetPhyAddr();
642- }642+ }
643- // y = cast((x * rstd * gamma + beta) * scales + zeropints)643+ // y = cast((x * rstd * gamma + beta) * scales + zeropints)
644- if (numQ == 1) {644+ if (numQ == 1) {
645- // scales + zeropints shape [1]645+ // scales + zeropints shape [1]
646- // rstd shape [A, 1]646+ // rstd shape [A, 1]
647- __VEC_SCOPE__647+ __VEC_SCOPE__
648- {648+ {
649- RegTensor<float> xReg, rstdReg, gammaReg, betaReg;649+ RegTensor<float> xReg, rstdReg, gammaReg, betaReg;
650- RegTensor<float> scales1Reg, zeroPoints1Reg, scales2Reg, zeroPoints2Reg;650+ RegTensor<float> scales1Reg, zeroPoints1Reg, scales2Reg, zeroPoints2Reg;
651- RegTensor<float> mul1Reg, mul2Reg;651+ RegTensor<float> mul1Reg, mul2Reg;
652- RegTensor<float> scales1ResultReg, scales2ResultReg;652+ RegTensor<float> scales1ResultReg, scales2ResultReg;
653- MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();653+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
654- MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();654+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
655- MaskReg mask4Int4 = CreateMask<float, MaskPattern::H>();655+ MaskReg mask4Int4 = CreateMask<float, MaskPattern::H>();
656- 656+ 
657- // ld scales and zeropoints657+ // ld scales and zeropoints
658- LoadScalarForDtypeTIn(scales1Addr, scales1Reg, pregFull, 0);658+ LoadScalarForDtypeTIn(scales1Addr, scales1Reg, pregFull, 0);
659- if constexpr (HAS_ZEROPINTS1) {659+ if constexpr (HAS_ZEROPINTS1) {
660- LoadScalarForDtypeTIn(zeroPoints1Addr, zeroPoints1Reg, pregFull, 0);660+ LoadScalarForDtypeTIn(zeroPoints1Addr, zeroPoints1Reg, pregFull, 0);
661- }661+ }
662- if constexpr (HAS_SCALES2) {662+ if constexpr (HAS_SCALES2) {
663- LoadScalarForDtypeTIn(scales2Addr, scales2Reg, pregFull, 0);663+ LoadScalarForDtypeTIn(scales2Addr, scales2Reg, pregFull, 0);
664- }664+ }
665- if constexpr (HAS_ZEROPINTS2) {665+ if constexpr (HAS_ZEROPINTS2) {
666- LoadScalarForDtypeTIn(zeroPoints2Addr, zeroPoints2Reg, pregFull, 0);666+ LoadScalarForDtypeTIn(zeroPoints2Addr, zeroPoints2Reg, pregFull, 0);
667- }667+ }
668- for (uint16_t i = 0; i < loopsA; i++) {668+ for (uint16_t i = 0; i < loopsA; i++) {
669- // ld rstd669+ // ld rstd
670- uint32_t sregElewiseNum = numR;670+ uint32_t sregElewiseNum = numR;
671- LoadScalarForDtypeTIn(rstdAddr, rstdReg, pregFull, i);671+ LoadScalarForDtypeTIn(rstdAddr, rstdReg, pregFull, i);
672- for (uint16_t j = 0; j < loopsR; j++) {672+ for (uint16_t j = 0; j < loopsR; j++) {
673- MaskReg pregCurLoop = UpdateMask<float>(sregElewiseNum);673+ MaskReg pregCurLoop = UpdateMask<float>(sregElewiseNum);
674- LoadTensorForDtypeTIn(xAddr, xReg, pregCurLoop, (i * sregxGammaBetaAlign + j * vectorLenB32));674+ LoadTensorForDtypeTIn(xAddr, xReg, pregCurLoop, (i * sregxGammaBetaAlign + j * vectorLenB32));
675- Mul(mul1Reg, xReg, rstdReg, pregCurLoop);675+ Mul(mul1Reg, xReg, rstdReg, pregCurLoop);
676- LoadTensorForDtypeTIn(gammaAddr, gammaReg, pregCurLoop, (j * vectorLenB32));676+ LoadTensorForDtypeTIn(gammaAddr, gammaReg, pregCurLoop, (j * vectorLenB32));
677- Mul(mul2Reg, gammaReg, mul1Reg, pregCurLoop);677+ Mul(mul2Reg, gammaReg, mul1Reg, pregCurLoop);
678- if constexpr (HAS_BETA) {678+ if constexpr (HAS_BETA) {
679- LoadTensorForDtypeTIn(betaAddr, betaReg, pregCurLoop, (j * vectorLenB32));679+ LoadTensorForDtypeTIn(betaAddr, betaReg, pregCurLoop, (j * vectorLenB32));
680- Add(mul2Reg, mul2Reg, betaReg, pregCurLoop);680+ Add(mul2Reg, mul2Reg, betaReg, pregCurLoop);
681- }681+ }
682- if constexpr (IS_SCALES_DIV) {682+ if constexpr (IS_SCALES_DIV) {
683- Div(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);683+ Div(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);
684- } else {684+ } else {
685- Mul(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);685+ Mul(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);
686- }686+ }
687- 687+ 
688- if constexpr (HAS_ZEROPINTS1) {688+ if constexpr (HAS_ZEROPINTS1) {
689- Add(scales1ResultReg, scales1ResultReg, zeroPoints1Reg, pregCurLoop);689+ Add(scales1ResultReg, scales1ResultReg, zeroPoints1Reg, pregCurLoop);
690- }690+ }
691- 691+ 
692- StoreTensorForDtypeTOut(y1Addr, scales1ResultReg, pregCurLoop, mask4Int4,692+ StoreTensorForDtypeTOut(y1Addr, scales1ResultReg, pregCurLoop, mask4Int4,
693- (i * sregyAlign + j * vectorLenB32));693+ (i * sregyAlign + j * vectorLenB32));
694- 694+ 
695- if constexpr ((HAS_ZEROPINTS2 || HAS_SCALES2)) {695+ if constexpr ((HAS_ZEROPINTS2 || HAS_SCALES2)) {
696- if constexpr (HAS_SCALES2) {696+ if constexpr (HAS_SCALES2) {
697- if constexpr (IS_SCALES_DIV) {697+ if constexpr (IS_SCALES_DIV) {
698- Div(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);698+ Div(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);
699- } else {699+ } else {
700- Mul(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);700+ Mul(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);
701- }701+ }
702- }702+ }
703- if constexpr (HAS_ZEROPINTS2) {703+ if constexpr (HAS_ZEROPINTS2) {
704- Add(scales2ResultReg, scales2ResultReg, zeroPoints2Reg, pregCurLoop);704+ Add(scales2ResultReg, scales2ResultReg, zeroPoints2Reg, pregCurLoop);
705- }705+ }
706- StoreTensorForDtypeTOut(y2Addr, scales2ResultReg, pregCurLoop, mask4Int4,706+ StoreTensorForDtypeTOut(y2Addr, scales2ResultReg, pregCurLoop, mask4Int4,
707- (i * sregyAlign + j * vectorLenB32));707+ (i * sregyAlign + j * vectorLenB32));
708- }708+ }
709- }709+ }
710- }710+ }
711- }711+ }
712- } else {712+ } else {
713- __VEC_SCOPE__713+ __VEC_SCOPE__
714- {714+ {
715- RegTensor<float> xReg, rstdReg, gammaReg, betaReg;715+ RegTensor<float> xReg, rstdReg, gammaReg, betaReg;
716- RegTensor<float> scales1Reg, zeroPoints1Reg, scales2Reg, zeroPoints2Reg;716+ RegTensor<float> scales1Reg, zeroPoints1Reg, scales2Reg, zeroPoints2Reg;
717- RegTensor<float> mul1Reg, mul2Reg;717+ RegTensor<float> mul1Reg, mul2Reg;
718- RegTensor<float> scales1ResultReg, scales2ResultReg;718+ RegTensor<float> scales1ResultReg, scales2ResultReg;
719- MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();719+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
720- MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();720+ MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
721- MaskReg mask4Int4 = CreateMask<float, MaskPattern::H>();721+ MaskReg mask4Int4 = CreateMask<float, MaskPattern::H>();
722- // ld scales and zeropoints722+ // ld scales and zeropoints
723- for (uint16_t i = 0; i < loopsA; i++) {723+ for (uint16_t i = 0; i < loopsA; i++) {
724- // ld rstd724+ // ld rstd
725- uint32_t sregElewiseNum = numR;725+ uint32_t sregElewiseNum = numR;
726- LoadScalarForDtypeTIn(rstdAddr, rstdReg, pregFull, i);726+ LoadScalarForDtypeTIn(rstdAddr, rstdReg, pregFull, i);
727- for (uint16_t j = 0; j < loopsR; j++) {727+ for (uint16_t j = 0; j < loopsR; j++) {
728- MaskReg pregCurLoop = UpdateMask<float>(sregElewiseNum);728+ MaskReg pregCurLoop = UpdateMask<float>(sregElewiseNum);
729- LoadTensorForDtypeTIn(xAddr, xReg, pregCurLoop, (i * sregxGammaBetaAlign + j * vectorLenB32));729+ LoadTensorForDtypeTIn(xAddr, xReg, pregCurLoop, (i * sregxGammaBetaAlign + j * vectorLenB32));
730- Mul(mul1Reg, xReg, rstdReg, pregCurLoop);730+ Mul(mul1Reg, xReg, rstdReg, pregCurLoop);
731- LoadTensorForDtypeTIn(gammaAddr, gammaReg, pregCurLoop, j * vectorLenB32);731+ LoadTensorForDtypeTIn(gammaAddr, gammaReg, pregCurLoop, j * vectorLenB32);
732- Mul(mul2Reg, gammaReg, mul1Reg, pregCurLoop);732+ Mul(mul2Reg, gammaReg, mul1Reg, pregCurLoop);
733- if constexpr (HAS_BETA) {733+ if constexpr (HAS_BETA) {
734- LoadTensorForDtypeTIn(betaAddr, betaReg, pregCurLoop, j * vectorLenB32);734+ LoadTensorForDtypeTIn(betaAddr, betaReg, pregCurLoop, j * vectorLenB32);
735- Add(mul2Reg, mul2Reg, betaReg, pregCurLoop);735+ Add(mul2Reg, mul2Reg, betaReg, pregCurLoop);
736- }736+ }
737- LoadTensorForDtypeTIn(scales1Addr, scales1Reg, pregCurLoop, j * vectorLenB32);737+ LoadTensorForDtypeTIn(scales1Addr, scales1Reg, pregCurLoop, j * vectorLenB32);
738- if constexpr (IS_SCALES_DIV) {738+ if constexpr (IS_SCALES_DIV) {
739- Div(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);739+ Div(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);
740- } else {740+ } else {
741- Mul(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);741+ Mul(scales1ResultReg, mul2Reg, scales1Reg, pregCurLoop);
742- }742+ }
743- 743+ 
744- if constexpr (HAS_ZEROPINTS1) {744+ if constexpr (HAS_ZEROPINTS1) {
745- LoadTensorForDtypeTIn(zeroPoints1Addr, zeroPoints1Reg, pregCurLoop, j * vectorLenB32);745+ LoadTensorForDtypeTIn(zeroPoints1Addr, zeroPoints1Reg, pregCurLoop, j * vectorLenB32);
746- Add(scales1ResultReg, scales1ResultReg, zeroPoints1Reg, pregCurLoop);746+ Add(scales1ResultReg, scales1ResultReg, zeroPoints1Reg, pregCurLoop);
747- }747+ }
748- 748+ 
749- StoreTensorForDtypeTOut(y1Addr, scales1ResultReg, pregCurLoop, mask4Int4,749+ StoreTensorForDtypeTOut(y1Addr, scales1ResultReg, pregCurLoop, mask4Int4,
750- (i * sregyAlign + j * vectorLenB32));750+ (i * sregyAlign + j * vectorLenB32));
751- 751+ 
752- if constexpr ((HAS_ZEROPINTS2 || HAS_SCALES2)) {752+ if constexpr ((HAS_ZEROPINTS2 || HAS_SCALES2)) {
753- if constexpr (HAS_SCALES2) {753+ if constexpr (HAS_SCALES2) {
754- LoadTensorForDtypeTIn(scales2Addr, scales2Reg, pregCurLoop, j * vectorLenB32);754+ LoadTensorForDtypeTIn(scales2Addr, scales2Reg, pregCurLoop, j * vectorLenB32);
755- if constexpr (IS_SCALES_DIV) {755+ if constexpr (IS_SCALES_DIV) {
756- Div(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);756+ Div(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);
757- } else {757+ } else {
758- Mul(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);758+ Mul(scales2ResultReg, mul2Reg, scales2Reg, pregCurLoop);
759- }759+ }
760- }760+ }
761- if constexpr (HAS_ZEROPINTS2) {761+ if constexpr (HAS_ZEROPINTS2) {
762- LoadTensorForDtypeTIn(zeroPoints2Addr, zeroPoints2Reg, pregCurLoop, j * vectorLenB32);762+ LoadTensorForDtypeTIn(zeroPoints2Addr, zeroPoints2Reg, pregCurLoop, j * vectorLenB32);
763- Add(scales2ResultReg, scales2ResultReg, zeroPoints2Reg, pregCurLoop);763+ Add(scales2ResultReg, scales2ResultReg, zeroPoints2Reg, pregCurLoop);
764- }764+ }
765- StoreTensorForDtypeTOut(y2Addr, scales2ResultReg, pregCurLoop, mask4Int4,765+ StoreTensorForDtypeTOut(y2Addr, scales2ResultReg, pregCurLoop, mask4Int4,
766- (i * sregyAlign + j * vectorLenB32));766+ (i * sregyAlign + j * vectorLenB32));
767- }767+ }
768- }768+ }
769- }769+ }
770- }770+ }
771- }771+ }
772- }772+ }
773-};773+};
774-} // namespace RmsNormQuantV2774+} // namespace RmsNormQuantV2
775-#endif // RMS_NORM_QUANT_V2_REBASE_H_775+#endif // RMS_NORM_QUANT_V2_REBASE_H_
Mnorm/rms_norm_quant_v2/op_kernel/arch35/rms_norm_quant_v2_regbase_recompute.h+39-40
@@ -377,15 +377,14 @@ public:
377 }377 }
378 378 
379 // 输出转连续 ub 到 ub 搬运379 // 输出转连续 ub 到 ub 搬运
380- __local_mem__ float* dstPtr = (__local_mem__ float*)rstdLocal.GetPhyAddr();380+ __ubuf__ float* dstPtr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
381- __local_mem__ float* cachePtr = (__local_mem__ float*)cacheLocal.GetPhyAddr() +381+ __ubuf__ float* cachePtr = (__ubuf__ float*)cacheLocal.GetPhyAddr() + resultCacheID_ * AR_RECOMPUTE_SUM_LEN;
382- resultCacheID_ * AR_RECOMPUTE_SUM_LEN;
383 __VEC_SCOPE__382 __VEC_SCOPE__
384 {383 {
385 RegTensor<float> a;384 RegTensor<float> a;
386 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();385 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
387- DataCopy<float, LoadDist::DIST_NORM>(a, cachePtr);386+ LoadAlign<float, LoadDist::DIST_NORM>(a, cachePtr);
388- DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstPtr + rowIndex, a, pregOne);387+ StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dstPtr + rowIndex, a, pregOne);
389 }388 }
390 }389 }
391 390 
@@ -394,21 +393,21 @@ public:
394 {393 {
395 uint16_t innerLoopTimes = cacheId;394 uint16_t innerLoopTimes = cacheId;
396 uint32_t innerLoopStride = stride;395 uint32_t innerLoopStride = stride;
397- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();396+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
398- __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride;397+ __ubuf__ float* cache = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheId * stride;
399- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();398+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
400 399 
401 __VEC_SCOPE__400 __VEC_SCOPE__
402 {401 {
403 RegTensor<float> aReg, bReg;402 RegTensor<float> aReg, bReg;
404 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();403 MaskReg pregOne = CreateMask<float, MaskPattern::VL1>();
405 404 
406- DataCopy(aReg, (__local_mem__ float*)src);405+ LoadAlign(aReg, (__ubuf__ float*)src);
407 for (uint16_t j = 0; j < innerLoopTimes; ++j) {406 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
408- DataCopy(bReg, dst + j * innerLoopStride);407+ LoadAlign(bReg, dst + j * innerLoopStride);
409 Add(aReg, aReg, bReg, pregOne);408 Add(aReg, aReg, bReg, pregOne);
410 }409 }
411- DataCopy((__local_mem__ float*)cache, aReg, pregOne);410+ StoreAlign((__ubuf__ float*)cache, aReg, pregOne);
412 }411 }
413 }412 }
414 413 
@@ -463,31 +462,31 @@ public:
463 SetOverflowMode<T_Y>(0);462 SetOverflowMode<T_Y>(0);
464 uint32_t sreg = (uint32_t)count;463 uint32_t sreg = (uint32_t)count;
465 uint16_t repeatTimes = CeilDivision(count, VL_FP32);464 uint16_t repeatTimes = CeilDivision(count, VL_FP32);
466- __local_mem__ T_X* xAddr = (__local_mem__ T_X*)xLocal.GetPhyAddr();465+ __ubuf__ T_X* xAddr = (__ubuf__ T_X*)xLocal.GetPhyAddr();
467- __local_mem__ float* rstdAddr = (__local_mem__ float*)rstdLocal.GetPhyAddr();466+ __ubuf__ float* rstdAddr = (__ubuf__ float*)rstdLocal.GetPhyAddr();
468- __local_mem__ T_X* gammaAddr = (__local_mem__ T_X*)gammaLocal.GetPhyAddr();467+ __ubuf__ T_X* gammaAddr = (__ubuf__ T_X*)gammaLocal.GetPhyAddr();
469 468 
470- __local_mem__ T_SCALES* scales1Addr = (__local_mem__ T_SCALES*)scales1Local.GetPhyAddr();469+ __ubuf__ T_SCALES* scales1Addr = (__ubuf__ T_SCALES*)scales1Local.GetPhyAddr();
471- __local_mem__ T_SCALES* scales2Addr;470+ __ubuf__ T_SCALES* scales2Addr;
472 if constexpr (HAS_SCALES2) {471 if constexpr (HAS_SCALES2) {
473- scales2Addr = (__local_mem__ T_SCALES*)scales2Local.GetPhyAddr();472+ scales2Addr = (__ubuf__ T_SCALES*)scales2Local.GetPhyAddr();
474 }473 }
475- __local_mem__ T_ZEROPOINTS* zeroPoints1Addr;474+ __ubuf__ T_ZEROPOINTS* zeroPoints1Addr;
476- __local_mem__ T_ZEROPOINTS* zeroPoints2Addr;475+ __ubuf__ T_ZEROPOINTS* zeroPoints2Addr;
477 if constexpr (HAS_ZEROPOINTS1) {476 if constexpr (HAS_ZEROPOINTS1) {
478- zeroPoints1Addr = (__local_mem__ T_ZEROPOINTS*)zeroPoints1Local.GetPhyAddr();477+ zeroPoints1Addr = (__ubuf__ T_ZEROPOINTS*)zeroPoints1Local.GetPhyAddr();
479 }478 }
480 if constexpr (HAS_ZEROPOINTS2) {479 if constexpr (HAS_ZEROPOINTS2) {
481- zeroPoints2Addr = (__local_mem__ T_ZEROPOINTS*)zeroPoints2Local.GetPhyAddr();480+ zeroPoints2Addr = (__ubuf__ T_ZEROPOINTS*)zeroPoints2Local.GetPhyAddr();
482 }481 }
483- __local_mem__ T_X* betaAddr;482+ __ubuf__ T_X* betaAddr;
484 if constexpr (HAS_BETA) {483 if constexpr (HAS_BETA) {
485- betaAddr = (__local_mem__ T_X*)betaLocal.GetPhyAddr();484+ betaAddr = (__ubuf__ T_X*)betaLocal.GetPhyAddr();
486 }485 }
487- __local_mem__ yCopyDtype* y1Addr = (__local_mem__ yCopyDtype*)y1Local.GetPhyAddr();486+ __ubuf__ yCopyDtype* y1Addr = (__ubuf__ yCopyDtype*)y1Local.GetPhyAddr();
488- __local_mem__ yCopyDtype* y2Addr;487+ __ubuf__ yCopyDtype* y2Addr;
489 if constexpr (HAS_SCALES2) {488 if constexpr (HAS_SCALES2) {
490- y2Addr = (__local_mem__ yCopyDtype*)y2Local.GetPhyAddr();489+ y2Addr = (__ubuf__ yCopyDtype*)y2Local.GetPhyAddr();
491 }490 }
492 491 
493 if constexpr (NEED_BRC) {492 if constexpr (NEED_BRC) {
@@ -510,7 +509,7 @@ public:
510 if constexpr (HAS_ZEROPOINTS2) {509 if constexpr (HAS_ZEROPOINTS2) {
511 LoadScalarForDtypeTIn<T_ZEROPOINTS>(zeroPoints2Addr, zeroPoints2RegFp32, pregFull, 0);510 LoadScalarForDtypeTIn<T_ZEROPOINTS>(zeroPoints2Addr, zeroPoints2RegFp32, pregFull, 0);
512 }511 }
513- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + rstdOffset);512+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + rstdOffset);
514 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {513 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
515 maskReg = UpdateMask<float>(sreg);514 maskReg = UpdateMask<float>(sreg);
516 LoadTensorForDtypeTIn<T_X>(xAddr, xRegFp32, maskReg, i * VL_FP32);515 LoadTensorForDtypeTIn<T_X>(xAddr, xRegFp32, maskReg, i * VL_FP32);
@@ -556,7 +555,7 @@ public:
556 RegTensor<float> y1Reg, y2Reg;555 RegTensor<float> y1Reg, y2Reg;
557 MaskReg maskReg;556 MaskReg maskReg;
558 MaskReg mask4Int4 = CreateMask<float, MaskPattern::H>();557 MaskReg mask4Int4 = CreateMask<float, MaskPattern::H>();
559- DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + rstdOffset);558+ LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr + rstdOffset);
560 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {559 for (uint16_t i = 0; i < (uint16_t)repeatTimes; i++) {
561 maskReg = UpdateMask<float>(sreg);560 maskReg = UpdateMask<float>(sreg);
562 LoadTensorForDtypeTIn<T_X>(xAddr, xRegFp32, maskReg, i * VL_FP32);561 LoadTensorForDtypeTIn<T_X>(xAddr, xRegFp32, maskReg, i * VL_FP32);
@@ -662,9 +661,9 @@ private:
662 __aicore__ inline void FoldBlockVF(LocalTensor<T_X>& xLocal, LocalTensor<T_X>& xFoldLocal,661 __aicore__ inline void FoldBlockVF(LocalTensor<T_X>& xLocal, LocalTensor<T_X>& xFoldLocal,
663 LocalTensor<float> xFp32Tmp, uint32_t tailCount, uint32_t count)662 LocalTensor<float> xFp32Tmp, uint32_t tailCount, uint32_t count)
664 {663 {
665- __local_mem__ T_X* xInUb = (__local_mem__ T_X*)xLocal.GetPhyAddr();664+ __ubuf__ T_X* xInUb = (__ubuf__ T_X*)xLocal.GetPhyAddr();
666- __local_mem__ float* xFp32TmpBuf = (__local_mem__ float*)xFp32Tmp.GetPhyAddr();665+ __ubuf__ float* xFp32TmpBuf = (__ubuf__ float*)xFp32Tmp.GetPhyAddr();
667- __local_mem__ T_X* xFoldInUb = (__local_mem__ T_X*)xFoldLocal.GetPhyAddr();666+ __ubuf__ T_X* xFoldInUb = (__ubuf__ T_X*)xFoldLocal.GetPhyAddr();
668 667 
669 uint16_t loops = (count + VL_FP32 - 1) / VL_FP32;668 uint16_t loops = (count + VL_FP32 - 1) / VL_FP32;
670 uint16_t tailLoops = (tailCount + VL_FP32 - 1) / VL_FP32;669 uint16_t tailLoops = (tailCount + VL_FP32 - 1) / VL_FP32;
@@ -683,13 +682,13 @@ private:
683 Mul(xFoldReg, xFoldReg, xFoldReg, pregLoop);682 Mul(xFoldReg, xFoldReg, xFoldReg, pregLoop);
684 Add(sum, xReg, xFoldReg, pregLoop);683 Add(sum, xReg, xFoldReg, pregLoop);
685 Select(sum, sum, xReg, pregLoop);684 Select(sum, sum, xReg, pregLoop);
686- DataCopy<float, StoreDist::DIST_NORM_B32>(xFp32TmpBuf + offset, sum, pregFull);685+ StoreAlign<float, StoreDist::DIST_NORM_B32>(xFp32TmpBuf + offset, sum, pregFull);
687 }686 }
688 for (uint16_t i = 0; i < static_cast<uint16_t>(loops - tailLoops); ++i) {687 for (uint16_t i = 0; i < static_cast<uint16_t>(loops - tailLoops); ++i) {
689 uint32_t offset = (i + tailLoops) * VL_FP32;688 uint32_t offset = (i + tailLoops) * VL_FP32;
690 LoadTensorForDtypeTIn<T_X>(xInUb, xReg, pregFull, offset);689 LoadTensorForDtypeTIn<T_X>(xInUb, xReg, pregFull, offset);
691 Mul(xReg, xReg, xReg, pregFull);690 Mul(xReg, xReg, xReg, pregFull);
692- DataCopy<float, StoreDist::DIST_NORM_B32>(xFp32TmpBuf + offset, xReg, pregFull);691+ StoreAlign<float, StoreDist::DIST_NORM_B32>(xFp32TmpBuf + offset, xReg, pregFull);
693 }692 }
694 }693 }
695 }694 }
@@ -703,9 +702,9 @@ private:
703 uint32_t outerLoopStride = VL_FP32;702 uint32_t outerLoopStride = VL_FP32;
704 uint32_t innerLoopStride = stride;703 uint32_t innerLoopStride = stride;
705 704 
706- __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr();705+ __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr();
707- __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride;706+ __ubuf__ float* cache = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheId * stride;
708- __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr();707+ __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr();
709 708 
710 __VEC_SCOPE__709 __VEC_SCOPE__
711 {710 {
@@ -714,13 +713,13 @@ private:
714 AscendC::MicroAPI::MaskReg pMask;713 AscendC::MicroAPI::MaskReg pMask;
715 for (uint16_t i = 0; i < outerLoopTimes; ++i) {714 for (uint16_t i = 0; i < outerLoopTimes; ++i) {
716 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);715 pMask = AscendC::MicroAPI::UpdateMask<float>(sreg);
717- AscendC::MicroAPI::DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride);716+ AscendC::MicroAPI::LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride);
718 for (uint16_t j = 0; j < innerLoopTimes; ++j) {717 for (uint16_t j = 0; j < innerLoopTimes; ++j) {
719- AscendC::MicroAPI::DataCopy(bReg,718+ AscendC::MicroAPI::LoadAlign(bReg,
720- (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride);719+ (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride);
721 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);720 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask);
722 }721 }
723- AscendC::MicroAPI::DataCopy((__local_mem__ float*)cache + i * outerLoopStride, aReg, pMask);722+ AscendC::MicroAPI::StoreAlign((__ubuf__ float*)cache + i * outerLoopStride, aReg, pMask);
724 }723 }
725 }724 }
726 }725 }