/** * Copyright (c) 2026 Huawei Technologies Co., Ltd. * This program is free software, you can redistribute it and/or modify it under the terms and conditions of * CANN Open Software License Agreement Version 2.0 (the "License"). * Please refer to the License for details. You may not use this file except in compliance with the License. * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. * See LICENSE in the root of the software repository for the full text of the License. */ /*! * \file hc_post_bfloat16.h * \brief */ #ifndef HC_POST_BFLOAT16_H #define HC_POST_BFLOAT16_H #include "kernel_operator.h" namespace HcPostRegBase { using namespace AscendC; template class HcPostRegBaseBfloat16 { public: __aicore__ inline HcPostRegBaseBfloat16() {}; __aicore__ inline void Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, GM_ADDR workspace, const HcPostTilingData *tilingData, TPipe *pipe); __aicore__ inline void Process(); __aicore__ inline void DataCopyInX(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); __aicore__ inline void DataCopyInPost(int64_t batchIndex); __aicore__ inline void DataCopyInResidual(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); __aicore__ inline void DataCopyInComb(int64_t batchIndex); __aicore__ inline void DataCopyOut(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset); __aicore__ inline void DoProcess(int64_t batchSize); __aicore__ inline void DoCompute(LocalTensor sumTempBuf, LocalTensor postUb, LocalTensor combUb, int64_t batchIndex, int64_t dOffset, int64_t dDealing); __aicore__ inline void DoMulAndAdd(LocalTensor xUb, LocalTensor postUb, LocalTensor residualUb, LocalTensor combUb, LocalTensor sumTempBuf, int64_t dOnceDealing); private: TPipe* pipe_; const HcPostTilingData* tiling_; constexpr static AscendC::MicroAPI::CastTrait castB16ToB32 = { AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN, AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN }; int32_t blkIdx_ = -1; int64_t batch_ = 0; int64_t hcParam_ = 0; int64_t dParam_ = 0; int64_t batchOneCoreTail_ = 0; int64_t batchOneCore_ = 0; int64_t isFrontCore_ = 0; int64_t dParamAlign_ = 0; int64_t dOnceDealing_ = 0; int64_t dLastDealing_ = 0; int64_t dSplitTime_ = 0; int64_t dParamOnceAlign_ = 0; static constexpr int32_t ONE_BLOCK_SIZE = 32; int32_t perBlock32 = ONE_BLOCK_SIZE / sizeof(float); GlobalTensor xGm_; GlobalTensor residualGm_; GlobalTensor postGm_; GlobalTensor combGm_; GlobalTensor yGm_; TQue xQue_; TQue residualQue_; TQue postQue_; TQue combQue_; TQue sumQue_; TBuf sumTempBuf_; }; template __aicore__ inline void HcPostRegBaseBfloat16::Init(GM_ADDR x, GM_ADDR residual, GM_ADDR post, GM_ADDR comb, GM_ADDR y, GM_ADDR workspace, const HcPostTilingData *tilingData, TPipe *pipe) { blkIdx_ = GetBlockIdx(); if (blkIdx_ >= tilingData->usedCoreNum) { return; } tiling_ = tilingData; pipe_ = pipe; hcParam_ = tilingData->hcParam; dParam_ = tilingData->dParam; batchOneCoreTail_ = tilingData->batchOneCoreTail; batchOneCore_ = tilingData->batchOneCore; isFrontCore_ = blkIdx_ < tilingData->frontCore; int64_t frontCore = tilingData->frontCore; dOnceDealing_ = tilingData->dOnceDealing; dLastDealing_ = tilingData->dLastDealing; dSplitTime_ = tilingData->dSplitTime; dParamAlign_ = (dParam_ + perBlock32 - 1) / perBlock32 * perBlock32; dParamOnceAlign_ = (dOnceDealing_ + perBlock32 - 1) / perBlock32 * perBlock32; int64_t xOffset = blkIdx_ * batchOneCore_ * dParam_; int64_t residualOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; int64_t postOffset = blkIdx_ * batchOneCore_ * hcParam_; int64_t combOffset = blkIdx_ * batchOneCore_ * hcParam_ * hcParam_; int64_t yOffset = blkIdx_ * batchOneCore_ * hcParam_ * dParam_; if (!isFrontCore_) { xOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * dParam_; residualOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; postOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_; combOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * hcParam_; yOffset = (blkIdx_ * batchOneCoreTail_ + frontCore) * hcParam_ * dParam_; } xGm_.SetGlobalBuffer((__gm__ T1 *)x + xOffset); residualGm_.SetGlobalBuffer((__gm__ T1 *)residual + residualOffset); postGm_.SetGlobalBuffer((__gm__ T2 *)post + postOffset); combGm_.SetGlobalBuffer((__gm__ T2 *)comb + combOffset); yGm_.SetGlobalBuffer((__gm__ T1 *)y + yOffset); pipe_->InitBuffer(xQue_, 2, dParamOnceAlign_ * sizeof(T1)); pipe_->InitBuffer(residualQue_, 2, hcParam_ * dParamOnceAlign_ * sizeof(T1)); pipe_->InitBuffer(postQue_, 2, hcParam_ * sizeof(T2)); pipe_->InitBuffer(combQue_, 2, hcParam_ * hcParam_ * sizeof(T2)); pipe_->InitBuffer(sumQue_, 2, hcParam_* dParamOnceAlign_ * sizeof(T1)); pipe_->InitBuffer(sumTempBuf_, hcParam_ * dParamOnceAlign_ * sizeof(float)); } template __aicore__ inline void HcPostRegBaseBfloat16::Process() { if (blkIdx_ >= tiling_->usedCoreNum) { return; } if (isFrontCore_) { DoProcess(tiling_->batchOneCore); } else { DoProcess(tiling_->batchOneCoreTail); } } template __aicore__ inline void HcPostRegBaseBfloat16::DataCopyInX(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) { LocalTensor xUb = xQue_.AllocTensor(); DataCopyExtParams copyParams; copyParams.blockCount = 1; copyParams.blockLen = dOnceDealing * sizeof(T1); copyParams.srcStride = 0; copyParams.dstStride = 0; DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; DataCopyPad(xUb, xGm_[batchIndex * dParam_ + dOffset], copyParams, dataCopyPadParams); xQue_.EnQue(xUb); } template __aicore__ inline void HcPostRegBaseBfloat16::DataCopyInPost(int64_t batchIndex) { LocalTensor postUb = postQue_.AllocTensor(); DataCopyExtParams copyParams; copyParams.blockCount = 1; copyParams.blockLen = hcParam_ * sizeof(T2); copyParams.srcStride = 0; copyParams.dstStride = 0; DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; DataCopyPad(postUb, postGm_[batchIndex * hcParam_], copyParams, dataCopyPadParams); postQue_.EnQue(postUb); } template __aicore__ inline void HcPostRegBaseBfloat16::DataCopyInResidual(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) { LocalTensor residualUb = residualQue_.AllocTensor(); DataCopyExtParams copyParams; copyParams.blockCount = hcParam_; copyParams.blockLen = dOnceDealing * sizeof(T1); copyParams.srcStride = (dParamAlign_ - dOnceDealing) * sizeof(T1); copyParams.dstStride = 0; DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; DataCopyPad(residualUb, residualGm_[batchIndex * hcParam_ * dParam_ + dOffset], copyParams, dataCopyPadParams); residualQue_.EnQue(residualUb); } template __aicore__ inline void HcPostRegBaseBfloat16::DataCopyInComb(int64_t batchIndex) { LocalTensor combUb = combQue_.AllocTensor(); DataCopyExtParams copyParams; copyParams.blockCount = 1; copyParams.blockLen = hcParam_ * hcParam_ * sizeof(T2); copyParams.srcStride = 0; copyParams.dstStride = 0; DataCopyPadExtParams dataCopyPadParams{false, 0, 0, 0}; DataCopyPad(combUb, combGm_[batchIndex * hcParam_ * hcParam_], copyParams, dataCopyPadParams); combQue_.EnQue(combUb); } template __aicore__ inline void HcPostRegBaseBfloat16::DataCopyOut(int64_t batchIndex, int64_t dOnceDealing, int64_t dOffset) { LocalTensor outBuf = sumQue_.DeQue(); DataCopyExtParams copyParams; copyParams.blockCount = hcParam_; copyParams.blockLen = dOnceDealing * sizeof(T1); copyParams.srcStride = 0; copyParams.dstStride = (dParamAlign_ - dOnceDealing) * sizeof(T1); AscendC::DataCopyPad(yGm_[batchIndex * hcParam_ * dParam_ + dOffset], outBuf, copyParams); sumQue_.FreeTensor(outBuf); } template __aicore__ inline void HcPostRegBaseBfloat16::DoMulAndAdd(LocalTensor xUb, LocalTensor postUb, LocalTensor residualUb, LocalTensor combUb, LocalTensor sumTempBuf, int64_t dOnceDealing) { uint16_t aTimes = hcParam_; uint32_t xDealNumAlign = (dOnceDealing + perBlock32 - 1) / perBlock32 * perBlock32; uint32_t vfLen = 256 / sizeof(float); uint16_t repeatTimes = dOnceDealing / vfLen; uint16_t hcTimes = hcParam_; uint32_t tailNum = dOnceDealing % vfLen; uint16_t tailLoopTimes = tailNum == 0 ? 0 : 1; auto residualAddr = (__ubuf__ T1*)residualUb.GetPhyAddr(); auto combAddr = (__ubuf__ T2*)combUb.GetPhyAddr(); auto sumAddr = (__ubuf__ float*)sumTempBuf.GetPhyAddr(); auto xAddr = (__ubuf__ T1*)xUb.GetPhyAddr(); auto postAddr = (__ubuf__ T2*)postUb.GetPhyAddr(); __VEC_SCOPE__ { uint32_t xDealNum = static_cast(hcParam_ * dOnceDealing); AscendC::MicroAPI::RegTensor xReg; AscendC::MicroAPI::RegTensor postReg; AscendC::MicroAPI::RegTensor xRegFloat; AscendC::MicroAPI::RegTensor postRegFloat; AscendC::MicroAPI::RegTensor residualReg0; AscendC::MicroAPI::RegTensor residualReg1; AscendC::MicroAPI::RegTensor residualReg2; AscendC::MicroAPI::RegTensor residualReg3; AscendC::MicroAPI::RegTensor combReg0; AscendC::MicroAPI::RegTensor combReg1; AscendC::MicroAPI::RegTensor combReg2; AscendC::MicroAPI::RegTensor combReg3; AscendC::MicroAPI::RegTensor residualRegFloat0; AscendC::MicroAPI::RegTensor residualRegFloat1; AscendC::MicroAPI::RegTensor residualRegFloat2; AscendC::MicroAPI::RegTensor residualRegFloat3; AscendC::MicroAPI::RegTensor combRegFloat0; AscendC::MicroAPI::RegTensor combRegFloat1; AscendC::MicroAPI::RegTensor combRegFloat2; AscendC::MicroAPI::RegTensor combRegFloat3; AscendC::MicroAPI::RegTensor sumRegFloat; AscendC::MicroAPI::RegTensor sumTempReg0; AscendC::MicroAPI::RegTensor sumTempReg1; AscendC::MicroAPI::RegTensor sumTempReg2; AscendC::MicroAPI::RegTensor sumTempReg3; AscendC::MicroAPI::MaskReg pMask; AscendC::MicroAPI::MaskReg pregMain = AscendC::MicroAPI::CreateMask(); for (uint16_t hcIndex = 0; hcIndex < hcTimes; hcIndex++) { pMask = AscendC::MicroAPI::UpdateMask(xDealNum); if constexpr (sizeof(T2) == 2) { AscendC::MicroAPI::DataCopy(combReg0, combAddr+hcIndex); AscendC::MicroAPI::DataCopy(combReg1, combAddr+hcParam_+hcIndex); AscendC::MicroAPI::DataCopy(combReg2, combAddr+2*hcParam_+hcIndex); AscendC::MicroAPI::DataCopy(combReg3, combAddr+3*hcParam_+hcIndex); AscendC::MicroAPI::DataCopy(postReg, postAddr + hcIndex); AscendC::MicroAPI::Cast(combRegFloat0, combReg0, pregMain); AscendC::MicroAPI::Cast(combRegFloat1, combReg1, pregMain); AscendC::MicroAPI::Cast(combRegFloat2, combReg2, pregMain); AscendC::MicroAPI::Cast(combRegFloat3, combReg3, pregMain); AscendC::MicroAPI::Cast(postRegFloat, postReg, pregMain); } else { AscendC::MicroAPI::DataCopy(combRegFloat0, combAddr+hcIndex); AscendC::MicroAPI::DataCopy(combRegFloat1, combAddr+hcParam_+hcIndex); AscendC::MicroAPI::DataCopy(combRegFloat2, combAddr+2*hcParam_+hcIndex); AscendC::MicroAPI::DataCopy(combRegFloat3, combAddr+3*hcParam_+hcIndex); AscendC::MicroAPI::DataCopy(postRegFloat, postAddr + hcIndex); } for (uint16_t j = 0; j < repeatTimes; j++) { AscendC::MicroAPI::DataCopy(xReg, xAddr+j*vfLen); AscendC::MicroAPI::DataCopy(residualReg0, residualAddr+j*vfLen); AscendC::MicroAPI::DataCopy(residualReg1, residualAddr+xDealNumAlign+j*vfLen); AscendC::MicroAPI::DataCopy(residualReg2, residualAddr+2*xDealNumAlign+j*vfLen); AscendC::MicroAPI::DataCopy(residualReg3, residualAddr+3*xDealNumAlign+j*vfLen); AscendC::MicroAPI::Cast(residualRegFloat0, residualReg0, pMask); AscendC::MicroAPI::Cast(residualRegFloat1, residualReg1, pMask); AscendC::MicroAPI::Cast(residualRegFloat2, residualReg2, pMask); AscendC::MicroAPI::Cast(residualRegFloat3, residualReg3, pMask); AscendC::MicroAPI::Cast(xRegFloat, xReg, pMask); AscendC::MicroAPI::Mul(sumTempReg0, residualRegFloat0, combRegFloat0, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat3, combRegFloat3, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat1, combRegFloat1, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat2, combRegFloat2, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, xRegFloat, postRegFloat, pMask); AscendC::MicroAPI::DataCopy(sumAddr+hcIndex*xDealNumAlign+j*vfLen, sumTempReg0, pMask); } for (uint16_t k = 0; k < tailLoopTimes; k++) { AscendC::MicroAPI::DataCopy(xReg, xAddr+repeatTimes*vfLen); AscendC::MicroAPI::DataCopy(residualReg0, residualAddr+repeatTimes*vfLen); AscendC::MicroAPI::DataCopy(residualReg1, residualAddr+xDealNumAlign+repeatTimes*vfLen); AscendC::MicroAPI::DataCopy(residualReg2, residualAddr+2*xDealNumAlign+repeatTimes*vfLen); AscendC::MicroAPI::DataCopy(residualReg3, residualAddr+3*xDealNumAlign+repeatTimes*vfLen); AscendC::MicroAPI::Cast(residualRegFloat0, residualReg0, pMask); AscendC::MicroAPI::Cast(residualRegFloat1, residualReg1, pMask); AscendC::MicroAPI::Cast(residualRegFloat2, residualReg2, pMask); AscendC::MicroAPI::Cast(residualRegFloat3, residualReg3, pMask); AscendC::MicroAPI::Cast(xRegFloat, xReg, pMask); AscendC::MicroAPI::Mul(sumTempReg0, residualRegFloat0, combRegFloat0, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat3, combRegFloat3, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat1, combRegFloat1, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, residualRegFloat2, combRegFloat2, pMask); AscendC::MicroAPI::MulAddDst(sumTempReg0, xRegFloat, postRegFloat, pMask); AscendC::MicroAPI::DataCopy(sumAddr+hcIndex*xDealNumAlign+repeatTimes*vfLen, sumTempReg0, pMask); } } } } template __aicore__ inline void HcPostRegBaseBfloat16::DoCompute(LocalTensor sumTempBuf, LocalTensor postUb, LocalTensor combUb, int64_t batchIndex, int64_t dOffset, int64_t dDealing) { DataCopyInX(batchIndex, dDealing, dOffset); LocalTensor xUb = xQue_.DeQue(); DataCopyInResidual(batchIndex, dDealing, dOffset); LocalTensor residualUb = residualQue_.DeQue(); DoMulAndAdd(xUb, postUb, residualUb, combUb, sumTempBuf, dDealing); LocalTensor sumUb = sumQue_.AllocTensor(); AscendC::Cast(sumUb, sumTempBuf, AscendC::RoundMode::CAST_RINT, hcParam_ * dOnceDealing_); sumQue_.EnQue(sumUb); DataCopyOut(batchIndex, dDealing, dOffset); residualQue_.FreeTensor(residualUb); xQue_.FreeTensor(xUb); } template __aicore__ inline void HcPostRegBaseBfloat16::DoProcess(int64_t batchSize) { LocalTensor sumTempBuf = sumTempBuf_.Get(); for (int64_t batchIndex = 0; batchIndex < batchSize; batchIndex++) { DataCopyInPost(batchIndex); LocalTensor postUb = postQue_.DeQue(); DataCopyInComb(batchIndex); LocalTensor combUb = combQue_.DeQue(); int64_t dOffset = 0; for (int64_t dIndex = 0; dIndex < dSplitTime_; dIndex++) { dOffset = dIndex*dOnceDealing_; DoCompute(sumTempBuf, postUb, combUb, batchIndex, dOffset, dOnceDealing_); } if (dLastDealing_ != 0) { dOffset = dSplitTime_ * dOnceDealing_; DoCompute(sumTempBuf, postUb, combUb, batchIndex, dOffset, dLastDealing_); } combQue_.FreeTensor(combUb); postQue_.FreeTensor(postUb); } } } #endif