EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms.cpp Normal file
View File

@@ -0,0 +1,50 @@
/**
* 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_pre_inv_rms_apt.cpp
* \brief hc_pre_inv_rms kernel file
*/
#include "hc_pre_inv_rms_full_load.h"
#include "hc_pre_inv_rms_full_load_large_d.h"
#if defined(__DAV_C310__)
#include "hc_pre_inv_rms_full_load_regbase.h"
using namespace HcPreInvRmsRegbase;
#endif
#include "kernel_operator.h"
using namespace AscendC;
using namespace HcPreInvRms;
using namespace HcPreInvRmsLargeD;
#define FULL_LOAD_TILING_KEY 1000
#define FULL_LOAD_LARGE_D_TILING_KEY 1001
#define REGBASE_FULL_LOAD_TILING_KEY 2000
extern "C" __global__ __aicore__ void hc_pre_inv_rms(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling)
{
TPipe pipe;
GET_TILING_DATA(tilingData, tiling);
if (TILING_KEY_IS(FULL_LOAD_TILING_KEY)) {
HcPreInvRmsFullLoad<DTYPE_X> op;
op.Init(x, y, workspace, &tilingData, &pipe);
op.Process();
} else if (TILING_KEY_IS(FULL_LOAD_LARGE_D_TILING_KEY)) {
HcPreInvRmsFullLoadLargeD<DTYPE_X> op;
op.Init(x, y, workspace, &tilingData, &pipe);
op.Process();
}
#if defined(__DAV_C310__)
else if (TILING_KEY_IS(REGBASE_FULL_LOAD_TILING_KEY)) {
HcPreInvRmsFullLoadRegbase<DTYPE_X> op;
op.Init(x, y, workspace, &tilingData, &pipe);
op.Process();
}
#endif
}

248
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load.h Normal file
View File

@@ -0,0 +1,248 @@
/**
* 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_pre_inv_rms.h
* \brief inv rms file
*/
#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_
#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_
#include "kernel_operator.h"
namespace HcPreInvRms {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t FLOAT_BTYPE_SIZE = 4;
constexpr uint32_t PER_REPEAT_LEN_B32 = 64;
constexpr uint32_t UB_BLOCK_SIZE = 32;
constexpr int32_t B16_TYPE_BYTE_SIZE = 2;
constexpr int32_t B32_TYPE_BYTE_SIZE = 4;
constexpr int32_t ONE_COUNT = 1;
constexpr int32_t FOUR_FOLD = 4;
constexpr int32_t DST_REP_STRIDE = 1;
constexpr int32_t SRC_BLK_STRIDE = 1;
constexpr int32_t SRC_REP_STRIDE = 8;
template <typename T>
class HcPreInvRmsFullLoad {
public:
__aicore__ inline HcPreInvRmsFullLoad() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void Compute(uint64_t curUbFactorA);
__aicore__ inline void ComputeB16(uint64_t curUbFactorA);
__aicore__ inline void ComputeB32(uint64_t curUbFactorA);
__aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, 1> inQueueX;
TQue<QuePosition::VECOUT, 1> outQueueY;
TBuf<TPosition::VECCALC> castBuf;
TBuf<TPosition::VECCALC> reduceBuf;
GlobalTensor<T> xGm;
GlobalTensor<float> yGm;
int64_t A; // 输入数据 A 轴大小
int64_t R; // 输入数据 R 轴大小
int64_t blockNumA; // 使用的核数
int64_t blockFactorA; // 每个核处理的A个数
int64_t blockTailFactorA; // 尾核处理的A个数
int64_t ubFactorA; // 每次ub循环处理的A个数
int32_t blockIdx_;
float epsilon; // 算子参数
uint32_t curBlockFactorA; // 当前核处理的A个数
uint32_t rAlign;
uint32_t rAlignB32;
uint32_t reduceBufNum;
};
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe)
{
A = tiling->A;
R = tiling->R;
blockNumA = tiling->blockNumA;
blockFactorA = tiling->blockFactorA;
blockTailFactorA = tiling->blockTailFactorA;
ubFactorA = tiling->ubFactorA;
epsilon = tiling->epsilon;
rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T));
rAlignB32 = ((R * FLOAT_BTYPE_SIZE + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / FLOAT_BTYPE_SIZE);
pipe_ = pipe;
blockIdx_ = GetBlockIdx();
if (blockIdx_ < blockNumA - 1) {
this->curBlockFactorA = this->blockFactorA;
} else if (blockIdx_ == blockNumA - 1) {
this->curBlockFactorA = this->blockTailFactorA;
} else {
return;
}
xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R);
yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA);
// pipe alloc memory to queue, the unit is Bytes
pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T));
pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE);
reduceBufNum = (rAlignB32 + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32;
pipe_->InitBuffer(reduceBuf, ubFactorA * reduceBufNum * FLOAT_BTYPE_SIZE);
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
pipe_->InitBuffer(castBuf, ubFactorA * rAlignB32 * FLOAT_BTYPE_SIZE);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::Process()
{
if (blockIdx_ >= blockNumA) {
return;
}
uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; // Ub循环次数
uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; // 最后一次Ub循环的A轴大小
uint64_t curUbFactorA = ubFactorA;
for (uint64_t idx = 0; idx < aUbLoopCount; idx++) {
if (idx == aUbLoopCount - 1) {
curUbFactorA = tailUbFactorA;
}
CopyIn(idx, curUbFactorA);
Compute(curUbFactorA);
CopyOut(idx, curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::CopyIn(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
int64_t xGmStartAddr = idx * R * ubFactorA;
DataCopyExtParams dataCopyParams{
static_cast<uint16_t>(curUbFactorA), static_cast<uint32_t>(R * sizeof(T)), 0, 0, 0};
DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams);
inQueueX.EnQue<T>(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::Compute(uint64_t curUbFactorA)
{
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
ComputeB16(curUbFactorA);
} else if constexpr (sizeof(T) == B32_TYPE_BYTE_SIZE) {
ComputeB32(curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::ComputeB16(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
LocalTensor<float> castLocal = castBuf.Get<float>();
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
int32_t perFoldElems = (rAlignB32 + FOUR_FOLD - 1) / FOUR_FOLD; // 4096
int32_t perFoldRepTime = (perFoldElems + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; // 64
AscendC::Cast(castLocal, xLocal, AscendC::RoundMode::CAST_NONE, R);
PipeBarrier<PIPE_V>();
AscendC::Mul(castLocal, castLocal, castLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
for (int j = 0; j < FOUR_FOLD; j++) {
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal[idx * reduceBufNum + j * perFoldRepTime], castLocal[idx * rAlignB32 + j * perFoldElems], PER_REPEAT_LEN_B32, perFoldRepTime,
DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal, reduceLocal, PER_REPEAT_LEN_B32, FOUR_FOLD, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
PipeBarrier<PIPE_V>();
WholeReduceSum(yLocal[idx], reduceLocal, FOUR_FOLD, 1, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::ComputeB32(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
int32_t perFoldElems = (rAlignB32 + FOUR_FOLD - 1) / FOUR_FOLD; // 4096
int32_t perFoldRepTime = (perFoldElems + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32; // 64
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
PipeBarrier<PIPE_V>();
AscendC::Mul(xLocal, xLocal, xLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
for (int j = 0; j < FOUR_FOLD; j++) {
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal[idx * reduceBufNum + j * perFoldRepTime], xLocal[idx * rAlignB32 + j * perFoldElems], PER_REPEAT_LEN_B32, perFoldRepTime,
DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
PipeBarrier<PIPE_V>();
WholeReduceSum(reduceLocal, reduceLocal, PER_REPEAT_LEN_B32, FOUR_FOLD, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
PipeBarrier<PIPE_V>();
WholeReduceSum(yLocal[idx], reduceLocal, FOUR_FOLD, 1, DST_REP_STRIDE, SRC_BLK_STRIDE, SRC_REP_STRIDE);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoad<T>::CopyOut(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<float> yLocal = outQueueY.DeQue<float>();
AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(curUbFactorA * sizeof(float)), 0, 0, 0};
DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams);
outQueueY.FreeTensor(yLocal);
}
} // namespace HcPreInvRms
#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_H_

262
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_large_d.h Normal file
View File

@@ -0,0 +1,262 @@
/**
* 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_pre_inv_rms_full_load_large_d.h
* \brief inv rms for large d (R=28672, d=7168)
*/
#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_
#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_
#include "kernel_operator.h"
namespace HcPreInvRmsLargeD {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 1;
constexpr int32_t FLOAT_BTYPE_SIZE = 4;
constexpr uint32_t PER_REPEAT_LEN_B32 = 64;
constexpr uint32_t UB_BLOCK_SIZE = 32;
constexpr int32_t B16_TYPE_BYTE_SIZE = 2;
constexpr int32_t B32_TYPE_BYTE_SIZE = 4;
constexpr int32_t HALf_INTERVAL = 2;
constexpr int32_t INDEX_TWO = 2;
constexpr int32_t INDEX_FOUR = 4;
constexpr int32_t INDEX_EIGHT = 8;
constexpr int32_t INDEX_SIXTEEN = 16;
__aicore__ inline int32_t FindPowerTwo(int32_t n)
{
n |= n >> 1;
n |= n >> INDEX_TWO;
n |= n >> INDEX_FOUR;
n |= n >> INDEX_EIGHT;
n |= n >> INDEX_SIXTEEN;
return (n + 1) >> 1;
}
__aicore__ inline void ReduceSumHalfInterval(
const LocalTensor<float>& dst_local, const LocalTensor<float>& src_local, int32_t count)
{
if (likely(count > PER_REPEAT_LEN_B32)) {
int32_t bodyCount = FindPowerTwo(count);
int32_t tailCount = count - bodyCount;
if (tailCount > 0) {
Add(src_local, src_local, src_local[bodyCount], tailCount);
PipeBarrier<PIPE_V>();
}
while (bodyCount > PER_REPEAT_LEN_B32) {
bodyCount = bodyCount / HALf_INTERVAL;
Add(src_local, src_local, src_local[bodyCount], bodyCount);
PipeBarrier<PIPE_V>();
}
AscendCUtils::SetMask<float>(PER_REPEAT_LEN_B32);
} else {
AscendCUtils::SetMask<float>(count);
}
#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220
if (g_coreType == AIV) {
WholeReduceSum<float, false>(dst_local, src_local, MASK_PLACEHOLDER, 1, 0, 1, 0);
}
#else
WholeReduceSum<float, false>(dst_local, src_local, MASK_PLACEHOLDER, 1, 1, 1, DEFAULT_REPEAT_STRIDE);
#endif
PipeBarrier<PIPE_V>();
}
template <typename T>
class HcPreInvRmsFullLoadLargeD {
public:
__aicore__ inline HcPreInvRmsFullLoadLargeD() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void Compute(uint64_t curUbFactorA);
__aicore__ inline void ComputeB16(uint64_t curUbFactorA);
__aicore__ inline void ComputeB32(uint64_t curUbFactorA);
__aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
TBuf<TPosition::VECCALC> castBuf;
TBuf<TPosition::VECCALC> reduceBuf;
GlobalTensor<T> xGm;
GlobalTensor<float> yGm;
int64_t A;
int64_t R;
int64_t blockNumA;
int64_t blockFactorA;
int64_t blockTailFactorA;
int64_t ubFactorA;
int32_t blockIdx_;
float epsilon;
uint32_t curBlockFactorA;
uint32_t rAlign;
uint32_t rAlignB32;
uint32_t reduceBufNum;
};
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe)
{
A = tiling->A;
R = tiling->R;
blockNumA = tiling->blockNumA;
blockFactorA = tiling->blockFactorA;
blockTailFactorA = tiling->blockTailFactorA;
ubFactorA = tiling->ubFactorA;
epsilon = tiling->epsilon;
rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T));
rAlignB32 = ((R * FLOAT_BTYPE_SIZE + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / FLOAT_BTYPE_SIZE);
pipe_ = pipe;
blockIdx_ = GetBlockIdx();
if (blockIdx_ < blockNumA - 1) {
this->curBlockFactorA = this->blockFactorA;
} else if (blockIdx_ == blockNumA - 1) {
this->curBlockFactorA = this->blockTailFactorA;
} else {
return;
}
xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R);
yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA);
pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T));
pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE);
reduceBufNum = (rAlignB32 + PER_REPEAT_LEN_B32 - 1) / PER_REPEAT_LEN_B32;
pipe_->InitBuffer(reduceBuf, ubFactorA * reduceBufNum * FLOAT_BTYPE_SIZE);
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
pipe_->InitBuffer(castBuf, ubFactorA * rAlignB32 * FLOAT_BTYPE_SIZE);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::Process()
{
if (blockIdx_ >= blockNumA) {
return;
}
uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA;
uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA;
uint64_t curUbFactorA = ubFactorA;
for (uint64_t idx = 0; idx < aUbLoopCount; idx++) {
if (idx == aUbLoopCount - 1) {
curUbFactorA = tailUbFactorA;
}
CopyIn(idx, curUbFactorA);
Compute(curUbFactorA);
CopyOut(idx, curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::CopyIn(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
int64_t xGmStartAddr = idx * R * ubFactorA;
DataCopyExtParams dataCopyParams{
static_cast<uint16_t>(curUbFactorA), static_cast<uint32_t>(R * sizeof(T)), 0, 0, 0};
DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams);
inQueueX.EnQue<T>(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::Compute(uint64_t curUbFactorA)
{
if constexpr (sizeof(T) == B16_TYPE_BYTE_SIZE) {
ComputeB16(curUbFactorA);
} else if constexpr (sizeof(T) == B32_TYPE_BYTE_SIZE) {
ComputeB32(curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::ComputeB16(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
LocalTensor<float> castLocal = castBuf.Get<float>();
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
AscendC::Cast(castLocal, xLocal, AscendC::RoundMode::CAST_NONE, R);
PipeBarrier<PIPE_V>();
AscendC::Mul(castLocal, castLocal, castLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
ReduceSumHalfInterval(yLocal[idx], castLocal[idx * rAlignB32], R);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::ComputeB32(uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
LocalTensor<float> reduceLocal = reduceBuf.Get<float>();
PipeBarrier<PIPE_V>();
AscendC::Mul(xLocal, xLocal, xLocal, R);
for (int idx = 0; idx < curUbFactorA; idx++) {
ReduceSumHalfInterval(yLocal[idx], xLocal[idx * rAlignB32], R);
}
float meanCof = 1.0f / R;
PipeBarrier<PIPE_V>();
AscendC::Muls(yLocal, yLocal, meanCof, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Adds(yLocal, yLocal, epsilon, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Duplicate(reduceLocal, 1.0f, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Sqrt(yLocal, yLocal, curUbFactorA);
PipeBarrier<PIPE_V>();
AscendC::Div(yLocal, reduceLocal, yLocal, curUbFactorA);
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadLargeD<T>::CopyOut(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<float> yLocal = outQueueY.DeQue<float>();
AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(curUbFactorA * sizeof(float)), 0, 0, 0};
DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams);
outQueueY.FreeTensor(yLocal);
}
} // namespace HcPreInvRmsLargeD
#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_LARGE_D_H_

303
csrc/moe/hc_pre_inv_rms/op_kernel/hc_pre_inv_rms_full_load_regbase.h Normal file
View File

@@ -0,0 +1,303 @@
/**
* 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_pre_inv_rms.h
* \brief inv rms file
*/
#ifndef ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_
#define ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_
#include "kernel_operator.h"
namespace HcPreInvRmsRegbase {
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
constexpr int32_t FLOAT_BTYPE_SIZE = 4;
constexpr uint32_t VF_LEN_B32 = 64;
constexpr uint32_t UB_BLOCK_SIZE = 32;
constexpr uint32_t FOLD_FOUR = 4;
constexpr AscendC::MicroAPI::CastTrait castTraitB162B32Even = {
AscendC::MicroAPI::RegLayout::ZERO,
AscendC::MicroAPI::SatMode::UNKNOWN,
AscendC::MicroAPI::MaskMergeMode::ZEROING,
AscendC::RoundMode::UNKNOWN,
};
template <typename T>
__aicore__ inline void LoadInputData(AscendC::MicroAPI::RegTensor<float>& dst, __local_mem__ T* src, AscendC::MicroAPI::MaskReg pregLoop, uint32_t srcOffset)
{
if constexpr (IsSameType<T, float>::value) {
DataCopy(dst, src + srcOffset);
} else if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) {
AscendC::MicroAPI::RegTensor<T> tmp;
DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(tmp, src + srcOffset);
Cast<float, T, castTraitB162B32Even>(dst, tmp, pregLoop);
}
}
template <typename T>
class HcPreInvRmsFullLoadRegbase {
public:
__aicore__ inline HcPreInvRmsFullLoadRegbase() {};
__aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe);
__aicore__ inline void Process();
__aicore__ inline void CopyIn(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void Compute(uint64_t idx, uint64_t curUbFactorA);
__aicore__ inline void ComputeFullLoadVF(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA);
__aicore__ inline void ComputeFullLoadVfPerf(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA);
__aicore__ inline void CopyOut(uint64_t idx, uint64_t curUbFactorA);
private:
TPipe* pipe_;
TQue<QuePosition::VECIN, 1> inQueueX;
TQue<QuePosition::VECOUT, 1> outQueueY;
GlobalTensor<T> xGm;
GlobalTensor<float> yGm;
int64_t A; // 输入数据 A 轴大小
int64_t R; // 输入数据 R 轴大小
int64_t blockNumA; // 使用的核数
int64_t blockFactorA; // 每个核处理的A个数
int64_t blockTailFactorA; // 尾核处理的A个数
int64_t ubFactorA; // 每次ub循环处理的A个数
int32_t blockIdx_;
float epsilon; // 算子参数
uint32_t curBlockFactorA; // 当前核处理的A个数
uint32_t rAlign;
};
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::Init(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, const HcPreInvRmsFullLoadTilingData* tiling, TPipe* pipe)
{
A = tiling->A;
R = tiling->R;
blockNumA = tiling->blockNumA;
blockFactorA = tiling->blockFactorA;
blockTailFactorA = tiling->blockTailFactorA;
ubFactorA = tiling->ubFactorA;
epsilon = tiling->epsilon;
rAlign = ((R * sizeof(T) + UB_BLOCK_SIZE - 1) / UB_BLOCK_SIZE) * (UB_BLOCK_SIZE / sizeof(T));
pipe_ = pipe;
blockIdx_ = GetBlockIdx();
if (blockIdx_ < blockNumA - 1) {
this->curBlockFactorA = this->blockFactorA;
} else if (blockIdx_ == blockNumA - 1) {
this->curBlockFactorA = this->blockTailFactorA;
} else {
return;
}
xGm.SetGlobalBuffer((__gm__ T*)x + blockIdx_ * blockFactorA * R, curBlockFactorA * R);
yGm.SetGlobalBuffer((__gm__ float*)y + blockIdx_ * blockFactorA, curBlockFactorA);
// pipe alloc memory to queue, the unit is Bytes
pipe_->InitBuffer(inQueueX, BUFFER_NUM, ubFactorA * rAlign * sizeof(T));
pipe_->InitBuffer(outQueueY, BUFFER_NUM, ubFactorA * FLOAT_BTYPE_SIZE);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::Process()
{
if (blockIdx_ >= blockNumA) {
return;
}
uint64_t aUbLoopCount = (curBlockFactorA + ubFactorA - 1) / ubFactorA; // Ub循环次数
uint64_t tailUbFactorA = curBlockFactorA - (aUbLoopCount - 1) * ubFactorA; // 最后一次Ub循环的A轴大小
uint64_t curUbFactorA = ubFactorA;
for (uint64_t idx = 0; idx < aUbLoopCount; idx++) {
if (idx == aUbLoopCount - 1) {
curUbFactorA = tailUbFactorA;
}
CopyIn(idx, curUbFactorA);
Compute(idx, curUbFactorA);
CopyOut(idx, curUbFactorA);
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::CopyIn(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
DataCopyPadExtParams<T> dataCopyPadParams{false, 0, 0, 0};
int64_t xGmStartAddr = idx * R * ubFactorA;
DataCopyExtParams dataCopyParams{
static_cast<uint16_t>(curUbFactorA), static_cast<uint32_t>(R * sizeof(T)), 0, 0, 0};
DataCopyPad(xLocal, xGm[xGmStartAddr], dataCopyParams, dataCopyPadParams);
inQueueX.EnQue<T>(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::Compute(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<T> xLocal = inQueueX.DeQue<T>();
LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();
if (R % 256 == 0) {
ComputeFullLoadVfPerf(yLocal, xLocal, rAlign, R, curUbFactorA);
} else {
ComputeFullLoadVF(yLocal, xLocal, rAlign, R, curUbFactorA);
}
outQueueY.EnQue<float>(yLocal);
inQueueX.FreeTensor(xLocal);
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::ComputeFullLoadVF(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA)
{
__ubuf__ T* xAddr = (__ubuf__ T*)xLocal.GetPhyAddr();
__ubuf__ float* yAddr = (__ubuf__ float*)yLocal.GetPhyAddr();
uint32_t vfLen = VF_LEN_B32;
uint16_t iLoopNum = curUbFactorA;
uint16_t needLoopNum = (rAlign + vfLen - 1) / vfLen; // 需要VF循环次数
uint16_t fourLoopNum = (needLoopNum + FOLD_FOUR - 1) / FOLD_FOUR; // 需要四循环次数
uint16_t formerFourLoopNum = fourLoopNum - 1;
uint16_t tailFourLoop = needLoopNum - formerFourLoopNum * FOLD_FOUR; // 最后一次四循环需处理VF数
uint16_t formerFourLoopElems = formerFourLoopNum * FOLD_FOUR * vfLen; // 除尾块外，处理元素个数
uint32_t tailFourLoopElems = rNum - formerFourLoopElems; // 最后一次四循环需处理实际元素个数
__VEC_SCOPE__
{
AscendC::MicroAPI::RegTensor<float> vregX1;
AscendC::MicroAPI::RegTensor<float> vregX2;
AscendC::MicroAPI::RegTensor<float> vregX3;
AscendC::MicroAPI::RegTensor<float> vregX4;
AscendC::MicroAPI::RegTensor<float> vregX;
AscendC::MicroAPI::RegTensor<float> vregSum;
AscendC::MicroAPI::RegTensor<float> vregR;
AscendC::MicroAPI::RegTensor<float> vregOne;
AscendC::MicroAPI::MaskReg preg;
AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
AscendC::MicroAPI::MaskReg pregOne = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::VL1>();
AscendC::MicroAPI::Duplicate(vregOne, 1.0f);
for (uint16_t i = 0; i < iLoopNum; i++) {
AscendC::MicroAPI::Duplicate(vregSum, 0.0f); // 用于累加的vreg
for (uint16_t j = 0; j < formerFourLoopNum; j++) {
uint32_t srcOffset1 = i * rAlign + 4 * j * vfLen;
LoadInputData<T>(vregX1, xAddr, pregAll, srcOffset1);
AscendC::MicroAPI::Mul(vregX1, vregX1, vregX1, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX1, pregAll);
uint32_t srcOffset2 = i * rAlign + (4 * j + 1) * vfLen;
LoadInputData<T>(vregX2, xAddr, pregAll, srcOffset2);
AscendC::MicroAPI::Mul(vregX2, vregX2, vregX2, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX2, pregAll);
uint32_t srcOffset3 = i * rAlign + (4 * j + 2) * vfLen;
LoadInputData<T>(vregX3, xAddr, pregAll, srcOffset3);
AscendC::MicroAPI::Mul(vregX3, vregX3, vregX3, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX3, pregAll);
uint32_t srcOffset4 = i * rAlign + (4 * j + 3) * vfLen;
LoadInputData<T>(vregX4, xAddr, pregAll, srcOffset4);
AscendC::MicroAPI::Mul(vregX4, vregX4, vregX4, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX4, pregAll);
}
tailFourLoopElems = rNum - formerFourLoopElems;
for (uint16_t j = 0; j < tailFourLoop; j++) {
preg = AscendC::MicroAPI::UpdateMask<float>(tailFourLoopElems);
uint32_t srcOffset = i * rAlign + formerFourLoopElems + j * vfLen;
LoadInputData<T>(vregX, xAddr, preg, srcOffset);
AscendC::MicroAPI::Mul(vregX, vregX, vregX, preg);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX, pregAll);
}
Reduce(vregSum, vregSum, pregAll);
AscendC::MicroAPI::Duplicate(vregR, (float)rNum);
AscendC::MicroAPI::Div(vregSum, vregSum, vregR, pregOne);
AscendC::MicroAPI::Adds(vregSum, vregSum, epsilon, pregOne);
AscendC::MicroAPI::Sqrt(vregSum, vregSum, pregOne);
AscendC::MicroAPI::Div(vregSum, vregOne, vregSum, pregOne);
AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(yAddr + i, vregSum, pregOne);
}
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::ComputeFullLoadVfPerf(LocalTensor<float>& yLocal, LocalTensor<T>& xLocal, uint32_t rAlign, uint32_t rNum, uint64_t curUbFactorA)
{
__ubuf__ T* xAddr = (__ubuf__ T*)xLocal.GetPhyAddr();
__ubuf__ float* yAddr = (__ubuf__ float*)yLocal.GetPhyAddr();
uint32_t vfLen = VF_LEN_B32;
uint16_t iLoopNum = curUbFactorA;
uint16_t jLoopNum = (rAlign / vfLen) / FOLD_FOUR;
__VEC_SCOPE__
{
AscendC::MicroAPI::RegTensor<float> vregX1;
AscendC::MicroAPI::RegTensor<float> vregX2;
AscendC::MicroAPI::RegTensor<float> vregX3;
AscendC::MicroAPI::RegTensor<float> vregX4;
AscendC::MicroAPI::RegTensor<float> vregSum;
AscendC::MicroAPI::RegTensor<float> vregR;
AscendC::MicroAPI::RegTensor<float> vregOne;
AscendC::MicroAPI::MaskReg pregAll = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
AscendC::MicroAPI::MaskReg pregOne = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::VL1>();
AscendC::MicroAPI::Duplicate(vregOne, 1.0f);
for (uint16_t i = 0; i < iLoopNum; i++) {
AscendC::MicroAPI::Duplicate(vregSum, 0.0f); // 用于累加的vreg
for (uint16_t j = 0; j < jLoopNum; j++) {
uint32_t srcOffset1 = i * rAlign + 4 * j * vfLen;
LoadInputData<T>(vregX1, xAddr, pregAll, srcOffset1);
AscendC::MicroAPI::Mul(vregX1, vregX1, vregX1, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX1, pregAll);
uint32_t srcOffset2 = i * rAlign + (4 * j + 1) * vfLen;
LoadInputData<T>(vregX2, xAddr, pregAll, srcOffset2);
AscendC::MicroAPI::Mul(vregX2, vregX2, vregX2, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX2, pregAll);
uint32_t srcOffset3 = i * rAlign + (4 * j + 2) * vfLen;
LoadInputData<T>(vregX3, xAddr, pregAll, srcOffset3);
AscendC::MicroAPI::Mul(vregX3, vregX3, vregX3, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX3, pregAll);
uint32_t srcOffset4 = i * rAlign + (4 * j + 3) * vfLen;
LoadInputData<T>(vregX4, xAddr, pregAll, srcOffset4);
AscendC::MicroAPI::Mul(vregX4, vregX4, vregX4, pregAll);
AscendC::MicroAPI::Add(vregSum, vregSum, vregX4, pregAll);
}
Reduce(vregSum, vregSum, pregAll);
AscendC::MicroAPI::Duplicate(vregR, (float)rNum);
AscendC::MicroAPI::Div(vregSum, vregSum, vregR, pregOne);
AscendC::MicroAPI::Adds(vregSum, vregSum, epsilon, pregOne);
AscendC::MicroAPI::Sqrt(vregSum, vregSum, pregOne);
AscendC::MicroAPI::Div(vregSum, vregOne, vregSum, pregOne);
AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(yAddr + i, vregSum, pregOne);
}
}
}
template <typename T>
__aicore__ inline void HcPreInvRmsFullLoadRegbase<T>::CopyOut(uint64_t idx, uint64_t curUbFactorA)
{
LocalTensor<float> yLocal = outQueueY.DeQue<float>();
AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(curUbFactorA * sizeof(float)), 0, 0, 0};
DataCopyPad(yGm[idx * ubFactorA], yLocal, copyParams);
outQueueY.FreeTensor(yLocal);
}
} // namespace HcPreInvRmsRegbase
#endif // ASCENDC_HC_PRE_INV_RMS_FULL_LOAD_REGBASE_H_