@@ -0,0 +1,676 @@
|
||||
/**
|
||||
* 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 lightning_indexer_kernel.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef LIGHTNING_INDEXER_KERNEL_H
|
||||
#define LIGHTNING_INDEXER_KERNEL_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "kernel_operator_list_tensor_intf.h"
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "lib/matrix/matmul/tiling.h"
|
||||
#include "../lightning_indexer_common.h"
|
||||
#include "lightning_indexer_service_vector.h"
|
||||
#include "lightning_indexer_service_cube.h"
|
||||
|
||||
namespace LIKernel {
|
||||
using namespace LICommon;
|
||||
using namespace matmul;
|
||||
using AscendC::CacheMode;
|
||||
using AscendC::CrossCoreSetFlag;
|
||||
using AscendC::CrossCoreWaitFlag;
|
||||
|
||||
// 由于S2循环前，RunInfo还没有赋值，使用TempLoopInfo临时存放B、N、S1轴相关的信息；同时减少重复计算
|
||||
struct TempLoopInfo {
|
||||
uint32_t bN2Idx = 0;
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t n2Idx = 0U;
|
||||
uint32_t gS1Idx = 0U;
|
||||
uint32_t gS1LoopEnd = 0U; // gS1方向循环的结束Idx
|
||||
uint32_t s2LoopEnd = 0U; // S2方向循环的结束Idx
|
||||
uint32_t actS1Size = 1U; // 当前Batch循环处理的S1轴的实际大小
|
||||
uint32_t actS2Size = 0U;
|
||||
uint32_t actS2SizeOrig = 0U; // 压缩前s2
|
||||
bool curActSeqLenIsZero = false;
|
||||
bool needDealActS1LessThanS1 = false; // S1的实际长度小于shape的S1长度时，是否需要清理输出
|
||||
uint32_t actMBaseSize = 0U; // m轴(gS1)方向实际大小
|
||||
uint32_t mBasicSizeTail = 0U; // gS1方向循环的尾基本块大小
|
||||
uint32_t s2BasicSizeTail = 0U; // S2方向循环的尾基本块大小
|
||||
};
|
||||
|
||||
template <typename LIT>
|
||||
class LightningIndexerKernel {
|
||||
public:
|
||||
__aicore__ inline LightningIndexerKernel(){};
|
||||
__aicore__ inline void Init(__gm__ uint8_t *query, __gm__ uint8_t *key, __gm__ uint8_t *weights,
|
||||
__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK,
|
||||
__gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices, __gm__ uint8_t *sparseValues,
|
||||
__gm__ uint8_t *workspace, const LITilingData *__restrict tiling, TPipe *tPipe);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
// =================================类型定义区=================================
|
||||
static constexpr bool DT_W_FLAG = LIT::weightsTypeFlag;
|
||||
using Q_T = typename LIT::queryType;
|
||||
using K_T = typename LIT::keyType;
|
||||
using OUT_T = typename LIT::outputType;
|
||||
static constexpr bool PAGE_ATTENTION = LIT::pageAttention;
|
||||
static constexpr LI_LAYOUT LAYOUT_T = LIT::layout;
|
||||
static constexpr LI_LAYOUT K_LAYOUT_T = LIT::keyLayout;
|
||||
using W_T =
|
||||
typename LightningIndexerTypeTraits<Q_T, typename std::conditional<DT_W_FLAG,
|
||||
float, void>::type>::weightsType;
|
||||
|
||||
LightningIndexerServiceCube<LIT> matmulService;
|
||||
LightningIndexerServiceVector<LIT> vectorService;
|
||||
|
||||
// =================================常量区=================================
|
||||
static constexpr uint32_t SYNC_C1_V1_FLAG = 4;
|
||||
static constexpr uint32_t SYNC_V1_C1_FLAG = 5;
|
||||
|
||||
static constexpr uint32_t M_BASE_SIZE = 256;
|
||||
static constexpr uint32_t S1_BASE_SIZE = 4;
|
||||
static constexpr uint32_t S1_BASE_SIZE_SMALL = 2;
|
||||
static constexpr uint32_t S2_BASE_SIZE = 128;
|
||||
static constexpr uint32_t HEAD_DIM = 128;
|
||||
static constexpr uint32_t K_HEAD_NUM = 1;
|
||||
static constexpr uint32_t GM_ALIGN_BYTES = 512;
|
||||
|
||||
static constexpr int64_t LD_PREFETCH_LEN = 2;
|
||||
|
||||
protected:
|
||||
TPipe *pipe = nullptr;
|
||||
|
||||
// offset
|
||||
uint64_t queryCoreOffset = 0ULL;
|
||||
uint64_t keyCoreOffset = 0ULL;
|
||||
uint64_t weightsCoreOffset = 0ULL;
|
||||
uint64_t indiceOutCoreOffset = 0ULL;
|
||||
uint64_t valueOutCoreOffset = 0ULL;
|
||||
// ================================Global Buffer区=================================
|
||||
GlobalTensor<Q_T> queryGm;
|
||||
GlobalTensor<K_T> keyGm;
|
||||
GlobalTensor<W_T> weightsGm;
|
||||
|
||||
GlobalTensor<int32_t> indiceOutGm;
|
||||
GlobalTensor<K_T> valueOutGm;
|
||||
GlobalTensor<int32_t> blockTableGm;
|
||||
|
||||
GlobalTensor<uint32_t> actualSeqLengthsGmQ;
|
||||
GlobalTensor<uint32_t> actualSeqLengthsGmKv;
|
||||
|
||||
// ================================类成员变量====================================
|
||||
// aic、aiv核信息
|
||||
uint32_t tmpBlockIdx = 0U;
|
||||
uint32_t aiCoreIdx = 0U;
|
||||
uint32_t usedCoreNum = 0U;
|
||||
|
||||
LICommon::ConstInfo constInfo{};
|
||||
TempLoopInfo tempLoopInfo{};
|
||||
LICommon::SplitCoreInfo splitCoreInfo{};
|
||||
|
||||
// ================================Init functions==================================
|
||||
__aicore__ inline void InitTilingData(const LITilingData *__restrict tilingData);
|
||||
__aicore__ inline void InitBuffers();
|
||||
__aicore__ inline void InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK);
|
||||
// ================================Split Core================================
|
||||
__aicore__ inline void SplitCore(uint32_t curCoreIdx, uint32_t &coreNum, LICommon::SplitCoreInfo &info);
|
||||
__aicore__ inline uint32_t GetTotalBaseBlockNum();
|
||||
__aicore__ inline uint32_t GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size, uint32_t actS2SizeOrig);
|
||||
// ================================Process functions================================
|
||||
__aicore__ inline void ProcessMain();
|
||||
__aicore__ inline void ProcessBaseBlock(uint32_t loop, uint64_t s2LoopIdx,
|
||||
LICommon::RunInfo runInfo);
|
||||
__aicore__ inline void ProcessInvalid();
|
||||
// ================================Params Calc=====================================
|
||||
__aicore__ inline void CalcGS1LoopParams(uint32_t bN2Idx);
|
||||
__aicore__ inline void GetBN2Idx(uint32_t bN2Idx);
|
||||
__aicore__ inline uint32_t GetActualSeqLen(uint32_t bIdx, uint32_t actualLenDims, bool isAccumSeq,
|
||||
GlobalTensor<uint32_t> &actualSeqLengthsGmKv, uint32_t defaultSeqLen);
|
||||
__aicore__ inline void GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size,
|
||||
uint32_t &actS2Size, uint32_t &actS2SizeOrig);
|
||||
__aicore__ inline void CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx);
|
||||
__aicore__ inline void CalcRunInfo(uint32_t loop, uint32_t s2LoopIdx, LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start);
|
||||
};
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::InitTilingData(const LITilingData *__restrict tilingData)
|
||||
{
|
||||
usedCoreNum = tilingData->usedCoreNum;
|
||||
constInfo.batchSize = tilingData->bSize;
|
||||
constInfo.qHeadNum = constInfo.gSize = tilingData->gSize;
|
||||
constInfo.kSeqSize = tilingData->s2Size;
|
||||
constInfo.qSeqSize = tilingData->s1Size;
|
||||
constInfo.attenMaskFlag = (tilingData->sparseMode == 3);
|
||||
constInfo.kCacheBlockSize = tilingData->blockSize;
|
||||
constInfo.maxBlockNumPerBatch = tilingData->maxBlockNumPerBatch;
|
||||
constInfo.sparseCount = tilingData->sparseCount;
|
||||
constInfo.outputLayout = LAYOUT_T; // 输出和输入形状一致
|
||||
if constexpr (std::is_same_v<K_T, float16_t>) {
|
||||
constInfo.INVALID_VAL = 0xFC00;
|
||||
}else {
|
||||
constInfo.INVALID_VAL = 0xFF80;
|
||||
}
|
||||
if constexpr (LAYOUT_T == LI_LAYOUT::TND) {
|
||||
constInfo.isAccumSeqS1 = true;
|
||||
}
|
||||
if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) {
|
||||
constInfo.isAccumSeqS2 = true;
|
||||
}
|
||||
|
||||
constInfo.kHeadNum = K_HEAD_NUM;
|
||||
constInfo.headDim = HEAD_DIM;
|
||||
|
||||
if (constInfo.sparseCount > 2048) {
|
||||
constInfo.mBaseSize = S1_BASE_SIZE_SMALL * constInfo.gSize;
|
||||
constInfo.s1BaseSize = S1_BASE_SIZE_SMALL;
|
||||
} else {
|
||||
constInfo.mBaseSize = S1_BASE_SIZE * constInfo.gSize;
|
||||
constInfo.s1BaseSize = S1_BASE_SIZE;
|
||||
}
|
||||
constInfo.s2BaseSize = S2_BASE_SIZE;
|
||||
constInfo.returnValueFlag = tilingData->returnValue;
|
||||
constInfo.splitMFlag = (constInfo.gSize == 64 && constInfo.sparseCount <= 2048);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::InitBuffers()
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
vectorService.InitBuffers(pipe);
|
||||
} else {
|
||||
matmulService.InitBuffers(pipe);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::InitActualSeqLen(__gm__ uint8_t *actualSeqLengthsQ,
|
||||
__gm__ uint8_t *actualSeqLengthsK)
|
||||
{
|
||||
if (actualSeqLengthsQ == nullptr) {
|
||||
constInfo.actualLenQDims = 0;
|
||||
} else {
|
||||
constInfo.actualLenQDims = constInfo.batchSize;
|
||||
actualSeqLengthsGmQ.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsQ, constInfo.actualLenQDims);
|
||||
}
|
||||
if (actualSeqLengthsK == nullptr) {
|
||||
constInfo.actualLenDims = 0;
|
||||
} else {
|
||||
constInfo.actualLenDims = constInfo.batchSize;
|
||||
actualSeqLengthsGmKv.SetGlobalBuffer((__gm__ uint32_t *)actualSeqLengthsK, constInfo.actualLenDims);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline uint32_t LightningIndexerKernel<LIT>::GetActualSeqLen(uint32_t bIdx,
|
||||
uint32_t actualLenDims, bool isAccumSeq,
|
||||
GlobalTensor<uint32_t> &actualSeqLengthsGmKv,
|
||||
uint32_t defaultSeqLen)
|
||||
{
|
||||
if (actualLenDims == 0) {
|
||||
return defaultSeqLen;
|
||||
} else if (isAccumSeq && bIdx > 0) {
|
||||
return actualSeqLengthsGmKv.GetValue(bIdx) - actualSeqLengthsGmKv.GetValue(bIdx - 1);
|
||||
} else {
|
||||
return actualSeqLengthsGmKv.GetValue(bIdx);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::GetS1S2ActualSeqLen(uint32_t bIdx, uint32_t &actS1Size,
|
||||
uint32_t &actS2Size, uint32_t &actS2SizeOrig)
|
||||
{
|
||||
actS1Size = GetActualSeqLen(bIdx, constInfo.actualLenQDims, constInfo.isAccumSeqS1, actualSeqLengthsGmQ,
|
||||
constInfo.qSeqSize);
|
||||
actS2SizeOrig =
|
||||
GetActualSeqLen(bIdx, constInfo.actualLenDims,
|
||||
constInfo.isAccumSeqS2, actualSeqLengthsGmKv, constInfo.kSeqSize);
|
||||
actS2Size = actS2SizeOrig;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline uint32_t LightningIndexerKernel<LIT>::GetS2BaseBlockNumOnMask(uint32_t s1gIdx, uint32_t actS1Size,
|
||||
uint32_t actS2SizeOrig)
|
||||
{
|
||||
if (actS2SizeOrig == 0) {
|
||||
return 0;
|
||||
}
|
||||
uint32_t s1Offset = constInfo.s1BaseSize * s1gIdx;
|
||||
int32_t validS2LenBase = static_cast<int32_t>(actS2SizeOrig) - static_cast<int32_t>(actS1Size);
|
||||
int32_t validS2Len = (static_cast<int32_t>(s1Offset) + validS2LenBase + static_cast<int32_t>(constInfo.s1BaseSize));
|
||||
validS2Len = Min(validS2Len, static_cast<int32_t>(actS2SizeOrig));
|
||||
validS2Len = Max(validS2Len, 1);
|
||||
return (validS2Len + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline uint32_t LightningIndexerKernel<LIT>::GetTotalBaseBlockNum()
|
||||
{
|
||||
uint32_t totalBlockNum = 0;
|
||||
uint32_t actS1Size, actS2Size, actS2SizeOrig;
|
||||
uint32_t s1GBaseNum, s2BaseNum;
|
||||
for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; bIdx++) {
|
||||
GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig);
|
||||
s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize);
|
||||
if (!constInfo.attenMaskFlag) {
|
||||
s2BaseNum = constInfo.isLDOpen ? CeilDiv(actS2Size, constInfo.s2BaseSize) : (actS2Size > 0 ? 1 : 0);
|
||||
totalBlockNum += s1GBaseNum * s2BaseNum * constInfo.kHeadNum;
|
||||
continue;
|
||||
}
|
||||
for (uint32_t s1gIdx = 0; s1gIdx < s1GBaseNum; s1gIdx++) {
|
||||
s2BaseNum = constInfo.isLDOpen
|
||||
? GetS2BaseBlockNumOnMask(s1gIdx, actS1Size, actS2SizeOrig)
|
||||
: (actS2Size > 0 ? 1 : 0);
|
||||
totalBlockNum += s2BaseNum * constInfo.kHeadNum;
|
||||
}
|
||||
}
|
||||
return totalBlockNum;
|
||||
}
|
||||
|
||||
|
||||
// 多核版本，双闭区间。基本原则：计算每个核最少处理的块数, 剩余的部分前面的核每个核多处理一块
|
||||
template <typename LIT>
|
||||
__aicore__ void inline LightningIndexerKernel<LIT>::SplitCore(uint32_t curCoreIdx, uint32_t &coreNum,
|
||||
LICommon::SplitCoreInfo &info)
|
||||
{
|
||||
uint32_t totalBlockNum = GetTotalBaseBlockNum();
|
||||
uint32_t minBlockPerCore = totalBlockNum / coreNum;
|
||||
uint32_t deal1MoreBlockCoreNum = totalBlockNum % coreNum;
|
||||
uint32_t coreIdx = 0;
|
||||
uint32_t lastGS1RemainBlockCnt = 0;
|
||||
uint32_t coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore;
|
||||
coreNum = minBlockPerCore == 0 ? deal1MoreBlockCoreNum : coreNum;
|
||||
if (curCoreIdx < coreNum) {
|
||||
splitCoreInfo.isCoreEnable = true;
|
||||
} else {
|
||||
splitCoreInfo.isCoreEnable = false;
|
||||
return;
|
||||
}
|
||||
|
||||
bool findLastCoreEnd = true;
|
||||
uint32_t actS1Size, actS2Size, actS2SizeOrig;
|
||||
uint32_t s1GBaseNum, s2BaseNum, s2Loop;
|
||||
for (uint32_t bN2Idx = 0; bN2Idx < constInfo.batchSize * constInfo.kHeadNum; bN2Idx++) {
|
||||
uint32_t bIdx = bN2Idx / constInfo.kHeadNum;
|
||||
if (bN2Idx % constInfo.kHeadNum == 0) {
|
||||
GetS1S2ActualSeqLen(bIdx, actS1Size, actS2Size, actS2SizeOrig);
|
||||
s1GBaseNum = CeilDiv(actS1Size, constInfo.s1BaseSize);
|
||||
s2BaseNum = CeilDiv(actS2Size, constInfo.s2BaseSize);
|
||||
}
|
||||
if constexpr (LAYOUT_T == LI_LAYOUT::BSND) {
|
||||
if (findLastCoreEnd && (s1GBaseNum == 0U || s2BaseNum == 0U)) {
|
||||
info.bN2Start = bN2Idx;
|
||||
info.gS1Start = 0;
|
||||
info.s2Start = 0;
|
||||
findLastCoreEnd = false;
|
||||
}
|
||||
}
|
||||
for (uint32_t gS1Idx = 0; gS1Idx < s1GBaseNum; gS1Idx++) {
|
||||
if (constInfo.attenMaskFlag) {
|
||||
s2BaseNum = GetS2BaseBlockNumOnMask(gS1Idx, actS1Size, actS2SizeOrig);
|
||||
}
|
||||
if (findLastCoreEnd && s2BaseNum == 0U) {
|
||||
info.bN2Start = bN2Idx;
|
||||
info.gS1Start = gS1Idx;
|
||||
info.s2Start = 0;
|
||||
findLastCoreEnd = false;
|
||||
}
|
||||
s2Loop = constInfo.isLDOpen ? s2BaseNum : (actS2Size > 0 ? 1 : 0);
|
||||
for (uint32_t s2Idx = 0; s2Idx < s2Loop;) {
|
||||
if (findLastCoreEnd) {
|
||||
info.bN2Start = bN2Idx;
|
||||
info.gS1Start = gS1Idx;
|
||||
info.s2Start = s2Idx;
|
||||
findLastCoreEnd = false;
|
||||
}
|
||||
uint32_t s2RemainBaseNum = s2Loop - s2Idx;
|
||||
if (lastGS1RemainBlockCnt + s2RemainBaseNum >= coreDealBlockCnt) {
|
||||
info.bN2End = bN2Idx;
|
||||
info.gS1End = gS1Idx;
|
||||
info.s2End = constInfo.isLDOpen
|
||||
? s2Idx + coreDealBlockCnt - lastGS1RemainBlockCnt - 1
|
||||
: s2BaseNum - 1;
|
||||
|
||||
if (coreIdx == curCoreIdx) {
|
||||
// S2被切N核，那么只有第一个核需要处理LD，其他核不用
|
||||
if (s2Idx == 0 && info.s2End + 1 < s2BaseNum) {
|
||||
info.isLD = true;
|
||||
}
|
||||
// 最后一个核处理的不是最后一个Batch，表明后面的Batch为空块(S2=0), 调整终点坐标以便清理输出
|
||||
if (coreIdx == coreNum - 1 && info.bN2End != constInfo.batchSize - 1) {
|
||||
info.bN2End = constInfo.batchSize - 1;
|
||||
info.gS1End = 0;
|
||||
info.s2End = 0;
|
||||
}
|
||||
return;
|
||||
}
|
||||
coreIdx++;
|
||||
findLastCoreEnd = true;
|
||||
s2Idx = info.s2End + 1;
|
||||
lastGS1RemainBlockCnt = 0;
|
||||
coreDealBlockCnt = coreIdx < deal1MoreBlockCoreNum ? minBlockPerCore + 1 : minBlockPerCore;
|
||||
} else {
|
||||
lastGS1RemainBlockCnt += s2RemainBaseNum;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::DealActSeqLenIsZero(uint32_t bIdx, uint32_t n2Idx, uint32_t s1Start)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
if (constInfo.outputLayout == LI_LAYOUT::TND) {
|
||||
uint32_t tSize = actualSeqLengthsGmQ.GetValue(constInfo.batchSize - 1);
|
||||
uint32_t tBase = bIdx == 0 ? 0 : actualSeqLengthsGmQ.GetValue(bIdx - 1);
|
||||
uint32_t s1Count = tempLoopInfo.actS1Size;
|
||||
|
||||
for (uint32_t s1Idx = s1Start; s1Idx < s1Count; s1Idx++) {
|
||||
uint64_t indiceOutOffset =
|
||||
(tBase + s1Idx) * constInfo.kHeadNum * constInfo.sparseCount + // T轴、s1轴偏移
|
||||
n2Idx * constInfo.sparseCount; // N2轴偏移
|
||||
vectorService.CleanInvalidOutput(indiceOutOffset);
|
||||
}
|
||||
} else if (constInfo.outputLayout == LI_LAYOUT::BSND) {
|
||||
for (uint32_t s1Idx = s1Start; s1Idx < constInfo.qSeqSize; s1Idx++) {
|
||||
// B,S1,N2,K
|
||||
uint64_t indiceOutOffset = bIdx * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount +
|
||||
s1Idx * constInfo.kHeadNum * constInfo.sparseCount + // B轴、S1轴偏移
|
||||
n2Idx * constInfo.sparseCount; // N2轴偏移
|
||||
vectorService.CleanInvalidOutput(indiceOutOffset);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::Init(__gm__ uint8_t *query,
|
||||
__gm__ uint8_t *key, __gm__ uint8_t *weights,
|
||||
__gm__ uint8_t *actualSeqLengthsQ, __gm__ uint8_t *actualSeqLengthsK,
|
||||
__gm__ uint8_t *blockTable, __gm__ uint8_t *sparseIndices,
|
||||
__gm__ uint8_t *sparseValues,
|
||||
__gm__ uint8_t *workspace, const LITilingData *__restrict tiling,
|
||||
TPipe *tPipe)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
tmpBlockIdx = GetBlockIdx(); // vec:0-47
|
||||
aiCoreIdx = tmpBlockIdx / 2;
|
||||
} else {
|
||||
tmpBlockIdx = GetBlockIdx(); // cube:0-23
|
||||
aiCoreIdx = tmpBlockIdx;
|
||||
}
|
||||
|
||||
InitTilingData(tiling);
|
||||
InitActualSeqLen(actualSeqLengthsQ, actualSeqLengthsK);
|
||||
|
||||
// 获取分核信息
|
||||
SplitCore(aiCoreIdx, usedCoreNum, splitCoreInfo);
|
||||
|
||||
pipe = tPipe;
|
||||
|
||||
uint64_t offset = 0;
|
||||
// vec 把整个s2的score存储在GM，大小为s1BaseSize * 16K * 4
|
||||
GlobalTensor<uint16_t> scoreGm; // 存放vec核写出的score
|
||||
uint64_t singleCoreScoreSize = constInfo.s1BaseSize *
|
||||
LICommon::Align(
|
||||
(uint64_t)constInfo.kSeqSize,
|
||||
(uint64_t)constInfo.s2BaseSize) *
|
||||
sizeof(uint16_t);
|
||||
scoreGm.SetGlobalBuffer((__gm__ uint16_t *)(workspace + aiCoreIdx * singleCoreScoreSize));
|
||||
offset += GetBlockNum() * singleCoreScoreSize;
|
||||
|
||||
if ASCEND_IS_AIV {
|
||||
vectorService.InitParams(constInfo, tiling);
|
||||
indiceOutGm.SetGlobalBuffer((__gm__ int32_t *)sparseIndices);
|
||||
valueOutGm.SetGlobalBuffer((__gm__ K_T *)sparseValues);
|
||||
weightsGm.SetGlobalBuffer((__gm__ W_T *)weights);
|
||||
blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
|
||||
vectorService.InitVecInputTensor(weightsGm, indiceOutGm, valueOutGm, blockTableGm);
|
||||
vectorService.InitVecWorkspaceTensor(scoreGm);
|
||||
} else {
|
||||
matmulService.InitParams(constInfo);
|
||||
queryGm.SetGlobalBuffer((__gm__ Q_T *)query);
|
||||
if constexpr (PAGE_ATTENTION) {
|
||||
blockTableGm.SetGlobalBuffer((__gm__ int32_t *)blockTable);
|
||||
}
|
||||
keyGm.SetGlobalBuffer((__gm__ K_T *)key);
|
||||
matmulService.InitMm1GlobalTensor(blockTableGm, keyGm, queryGm);
|
||||
}
|
||||
InitBuffers();
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::GetBN2Idx(uint32_t bN2Idx)
|
||||
{
|
||||
tempLoopInfo.bN2Idx = bN2Idx;
|
||||
tempLoopInfo.bIdx = bN2Idx / constInfo.kHeadNum;
|
||||
tempLoopInfo.n2Idx = bN2Idx % constInfo.kHeadNum;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::CalcS2LoopParams(uint32_t bN2LoopIdx, uint32_t gS1LoopIdx)
|
||||
{
|
||||
tempLoopInfo.gS1Idx = gS1LoopIdx;
|
||||
tempLoopInfo.actMBaseSize = constInfo.mBaseSize;
|
||||
uint32_t remainedGS1Size = tempLoopInfo.actS1Size * constInfo.gSize - tempLoopInfo.gS1Idx * constInfo.mBaseSize;
|
||||
if (remainedGS1Size <= constInfo.mBaseSize && remainedGS1Size > 0) {
|
||||
tempLoopInfo.actMBaseSize = tempLoopInfo.mBasicSizeTail;
|
||||
}
|
||||
|
||||
bool isEnd = (bN2LoopIdx == splitCoreInfo.bN2End) && (gS1LoopIdx == splitCoreInfo.gS1End);
|
||||
uint32_t s2BlockNum;
|
||||
if (constInfo.attenMaskFlag) {
|
||||
s2BlockNum = GetS2BaseBlockNumOnMask(gS1LoopIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2SizeOrig);
|
||||
} else {
|
||||
s2BlockNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
|
||||
}
|
||||
tempLoopInfo.s2LoopEnd = isEnd ? splitCoreInfo.s2End : s2BlockNum - 1;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::CalcGS1LoopParams(uint32_t bN2LoopIdx)
|
||||
{
|
||||
GetBN2Idx(bN2LoopIdx);
|
||||
GetS1S2ActualSeqLen(tempLoopInfo.bIdx, tempLoopInfo.actS1Size, tempLoopInfo.actS2Size, tempLoopInfo.actS2SizeOrig);
|
||||
if ((tempLoopInfo.actS2Size == 0) || (tempLoopInfo.actS1Size == 0)) {
|
||||
tempLoopInfo.curActSeqLenIsZero = true;
|
||||
return;
|
||||
}
|
||||
tempLoopInfo.curActSeqLenIsZero = false;
|
||||
tempLoopInfo.s2BasicSizeTail = tempLoopInfo.actS2Size % constInfo.s2BaseSize;
|
||||
tempLoopInfo.s2BasicSizeTail =
|
||||
(tempLoopInfo.s2BasicSizeTail == 0) ? constInfo.s2BaseSize : tempLoopInfo.s2BasicSizeTail;
|
||||
tempLoopInfo.mBasicSizeTail = (tempLoopInfo.actS1Size * constInfo.gSize) % constInfo.mBaseSize;
|
||||
tempLoopInfo.mBasicSizeTail =
|
||||
(tempLoopInfo.mBasicSizeTail == 0) ? constInfo.mBaseSize : tempLoopInfo.mBasicSizeTail;
|
||||
|
||||
uint32_t gS1SplitNum = (tempLoopInfo.actS1Size * constInfo.gSize + constInfo.mBaseSize - 1) / constInfo.mBaseSize;
|
||||
tempLoopInfo.gS1LoopEnd = (bN2LoopIdx == splitCoreInfo.bN2End) ? splitCoreInfo.gS1End : gS1SplitNum - 1;
|
||||
if constexpr (LAYOUT_T == LI_LAYOUT::BSND) {
|
||||
if (tempLoopInfo.gS1LoopEnd == gS1SplitNum - 1 && constInfo.qSeqSize > tempLoopInfo.actS1Size) {
|
||||
tempLoopInfo.needDealActS1LessThanS1 = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::CalcRunInfo(uint32_t loop,
|
||||
uint32_t s2LoopIdx, LICommon::RunInfo &runInfo)
|
||||
{
|
||||
runInfo.loop = loop;
|
||||
runInfo.bIdx = tempLoopInfo.bIdx;
|
||||
runInfo.gS1Idx = tempLoopInfo.gS1Idx;
|
||||
runInfo.s2Idx = s2LoopIdx;
|
||||
runInfo.bN2Idx = tempLoopInfo.bN2Idx;
|
||||
runInfo.isValid = s2LoopIdx <= tempLoopInfo.s2LoopEnd;
|
||||
|
||||
if (!runInfo.isValid) {
|
||||
return;
|
||||
}
|
||||
|
||||
runInfo.actS1Size = tempLoopInfo.actS1Size;
|
||||
runInfo.actS2Size = tempLoopInfo.actS2Size;
|
||||
runInfo.actS2SizeOrig = tempLoopInfo.actS2SizeOrig;
|
||||
// 计算实际基本块size
|
||||
runInfo.actMBaseSize = tempLoopInfo.actMBaseSize;
|
||||
runInfo.actualSingleProcessSInnerSize = constInfo.s2BaseSize;
|
||||
uint32_t s2SplitNum = (tempLoopInfo.actS2Size + constInfo.s2BaseSize - 1) / constInfo.s2BaseSize;
|
||||
if (runInfo.s2Idx == s2SplitNum - 1) {
|
||||
runInfo.actualSingleProcessSInnerSize = tempLoopInfo.s2BasicSizeTail;
|
||||
}
|
||||
runInfo.actualSingleProcessSInnerSizeAlign =
|
||||
LICommon::Align((uint32_t)runInfo.actualSingleProcessSInnerSize, LICommon::ConstInfo::BUFFER_SIZE_BYTE_32B);
|
||||
|
||||
runInfo.isFirstS2InnerLoop = s2LoopIdx == splitCoreInfo.s2Start;
|
||||
runInfo.isLastS2InnerLoop = s2LoopIdx == tempLoopInfo.s2LoopEnd;
|
||||
runInfo.isAllLoopEnd = (runInfo.bN2Idx == splitCoreInfo.bN2End) && (runInfo.gS1Idx == splitCoreInfo.gS1End) &&
|
||||
(runInfo.s2Idx == splitCoreInfo.s2End);
|
||||
|
||||
if (runInfo.isFirstS2InnerLoop) {
|
||||
uint64_t actualSeqQPrefixSum;
|
||||
if constexpr (LAYOUT_T == LI_LAYOUT::TND) {
|
||||
actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmQ.GetValue(runInfo.bIdx - 1);
|
||||
} else { // BSND
|
||||
actualSeqQPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.qSeqSize;
|
||||
}
|
||||
uint64_t tndBIdxOffset = actualSeqQPrefixSum * constInfo.qHeadNum * constInfo.headDim;
|
||||
// B,S1,N1(N2,G),D
|
||||
queryCoreOffset = tndBIdxOffset + runInfo.gS1Idx * constInfo.mBaseSize * constInfo.headDim;
|
||||
// B,S1,N1(N2,G)/T,N1(N2,G)
|
||||
weightsCoreOffset = actualSeqQPrefixSum * constInfo.qHeadNum + runInfo.n2Idx * constInfo.gSize;
|
||||
// B,S1,N2,k/T,N2,k
|
||||
indiceOutCoreOffset =
|
||||
actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount;
|
||||
// B,S1,N2,k/T,N2,k
|
||||
valueOutCoreOffset =
|
||||
actualSeqQPrefixSum * constInfo.kHeadNum * constInfo.sparseCount + runInfo.n2Idx * constInfo.sparseCount;
|
||||
}
|
||||
uint64_t actualSeqKPrefixSum;
|
||||
if constexpr (K_LAYOUT_T == LI_LAYOUT::TND) { // T N2 D
|
||||
actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : actualSeqLengthsGmKv.GetValue(runInfo.bIdx - 1);
|
||||
} else {
|
||||
actualSeqKPrefixSum = (runInfo.bIdx <= 0) ? 0 : runInfo.bIdx * constInfo.kSeqSize;
|
||||
}
|
||||
uint64_t tndBIdxOffsetForK = actualSeqKPrefixSum * constInfo.kHeadNum * constInfo.headDim;
|
||||
keyCoreOffset = tndBIdxOffsetForK + runInfo.s2Idx * constInfo.s2BaseSize * constInfo.kHeadNum * constInfo.headDim;
|
||||
runInfo.tensorQueryOffset = queryCoreOffset;
|
||||
runInfo.tensorKeyOffset = keyCoreOffset;
|
||||
runInfo.tensorWeightsOffset = weightsCoreOffset;
|
||||
runInfo.indiceOutOffset = indiceOutCoreOffset;
|
||||
runInfo.valueOutOffset = valueOutCoreOffset;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::Process()
|
||||
{
|
||||
if (usedCoreNum == 0) {
|
||||
// 没有计算任务，直接清理输出
|
||||
ProcessInvalid();
|
||||
return;
|
||||
}
|
||||
|
||||
ProcessMain();
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::ProcessInvalid()
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
uint32_t aivCoreNum = GetBlockNum() * 2; // 2 means c:v = 1:2
|
||||
uint64_t totalOutputSize =
|
||||
constInfo.batchSize * constInfo.qSeqSize * constInfo.kHeadNum * constInfo.sparseCount;
|
||||
uint64_t singleCoreSize =
|
||||
LICommon::Align((totalOutputSize + aivCoreNum - 1) / aivCoreNum, GM_ALIGN_BYTES / sizeof(OUT_T));
|
||||
uint64_t baseSize = tmpBlockIdx * singleCoreSize;
|
||||
if (baseSize < totalOutputSize) {
|
||||
uint64_t dealSize =
|
||||
(baseSize + singleCoreSize <= totalOutputSize) ? singleCoreSize : totalOutputSize - baseSize;
|
||||
GlobalTensor<OUT_T> output = indiceOutGm[baseSize];
|
||||
AscendC::InitGlobalMemory(output, dealSize, constInfo.INVALID_IDX);
|
||||
if (constInfo.returnValueFlag) {
|
||||
GlobalTensor<uint16_t> valueOutGmTmp;
|
||||
valueOutGmTmp.SetGlobalBuffer((__gm__ uint16_t *)valueOutGm.GetPhyAddr());
|
||||
GlobalTensor<uint16_t> valueOut = valueOutGmTmp[baseSize];
|
||||
AscendC::InitGlobalMemory(valueOut, dealSize, constInfo.INVALID_VAL);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::ProcessMain()
|
||||
{
|
||||
if (!splitCoreInfo.isCoreEnable) {
|
||||
return;
|
||||
}
|
||||
|
||||
if ASCEND_IS_AIV {
|
||||
vectorService.AllocEventID();
|
||||
CrossCoreSetFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_V>(LICommon::ConstInfo::CROSS_VC_EVENT + 0);
|
||||
CrossCoreSetFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_V>(LICommon::ConstInfo::CROSS_VC_EVENT + 1);
|
||||
} else {
|
||||
matmulService.AllocEventID();
|
||||
}
|
||||
|
||||
LICommon::RunInfo runInfo;
|
||||
uint32_t gloop = 0;
|
||||
for (uint32_t bN2LoopIdx = splitCoreInfo.bN2Start; bN2LoopIdx <= splitCoreInfo.bN2End; bN2LoopIdx++) {
|
||||
CalcGS1LoopParams(bN2LoopIdx);
|
||||
if (tempLoopInfo.curActSeqLenIsZero) {
|
||||
DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, 0U);
|
||||
continue;
|
||||
}
|
||||
for (uint32_t gS1LoopIdx = splitCoreInfo.gS1Start; gS1LoopIdx <= tempLoopInfo.gS1LoopEnd; gS1LoopIdx++) {
|
||||
CalcS2LoopParams(bN2LoopIdx, gS1LoopIdx);
|
||||
for (int s2LoopIdx = splitCoreInfo.s2Start; s2LoopIdx <= tempLoopInfo.s2LoopEnd; s2LoopIdx++) {
|
||||
ProcessBaseBlock(gloop, s2LoopIdx, runInfo);
|
||||
++gloop;
|
||||
}
|
||||
splitCoreInfo.s2Start = 0;
|
||||
}
|
||||
if (tempLoopInfo.needDealActS1LessThanS1) {
|
||||
DealActSeqLenIsZero(tempLoopInfo.bIdx, tempLoopInfo.n2Idx, tempLoopInfo.actS1Size);
|
||||
}
|
||||
splitCoreInfo.gS1Start = 0;
|
||||
}
|
||||
|
||||
if ASCEND_IS_AIV {
|
||||
vectorService.FreeEventID();
|
||||
} else {
|
||||
matmulService.FreeEventID();
|
||||
CrossCoreWaitFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_FIX>(LICommon::ConstInfo::CROSS_VC_EVENT + 0);
|
||||
CrossCoreWaitFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_FIX>(LICommon::ConstInfo::CROSS_VC_EVENT + 1);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerKernel<LIT>::ProcessBaseBlock(uint32_t loop,
|
||||
uint64_t s2LoopIdx, LICommon::RunInfo runInfo)
|
||||
{
|
||||
CalcRunInfo(loop, s2LoopIdx, runInfo);
|
||||
if ASCEND_IS_AIC {
|
||||
matmulService.ComputeMm1(runInfo);
|
||||
} else {
|
||||
vectorService.ProcessVec1(runInfo);
|
||||
if (runInfo.isLastS2InnerLoop) { // 本核s2last
|
||||
vectorService.ProcessTopK(runInfo);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace LIKernel
|
||||
#endif // LIGHTNING_INDEXER_KERNEL_H
|
||||
@@ -0,0 +1,470 @@
|
||||
/**
|
||||
* 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 lightning_indexer_service_cube.h
|
||||
* \brief use 5 buffer for matmul l1, better pipeline
|
||||
*/
|
||||
#ifndef LIGHTNING_INDEXER_SERVICE_CUBE_H
|
||||
#define LIGHTNING_INDEXER_SERVICE_CUBE_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "kernel_operator_list_tensor_intf.h"
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "lib/matrix/matmul/tiling.h"
|
||||
#include "../lightning_indexer_common.h"
|
||||
|
||||
namespace LIKernel {
|
||||
using namespace LICommon;
|
||||
template <typename LIT>
|
||||
class LightningIndexerServiceCube {
|
||||
public:
|
||||
using Q_T = typename LIT::queryType;
|
||||
using K_T = typename LIT::keyType;
|
||||
|
||||
__aicore__ inline LightningIndexerServiceCube(){};
|
||||
__aicore__ inline void InitBuffers(TPipe *pipe);
|
||||
__aicore__ inline void InitMm1GlobalTensor(const GlobalTensor<int32_t> &blkTableGm, const GlobalTensor<K_T> &keyGm,
|
||||
const GlobalTensor<Q_T> &queryGm);
|
||||
__aicore__ inline void InitParams(const ConstInfo &constInfo);
|
||||
__aicore__ inline void AllocEventID();
|
||||
__aicore__ inline void FreeEventID();
|
||||
__aicore__ inline void ComputeMm1(const LICommon::RunInfo &runInfo);
|
||||
|
||||
static constexpr uint64_t KEY_BUF_NUM = 3;
|
||||
static constexpr uint64_t QUERY_BUF_NUM = 2;
|
||||
static constexpr uint64_t L0_BUF_NUM = 2;
|
||||
|
||||
static constexpr uint32_t KEY_MTE1_MTE2_EVENT = EVENT_ID2;
|
||||
static constexpr uint32_t QUERY_MTE1_MTE2_EVENT = EVENT_ID5; // KEY_MTE1_MTE2_EVENT + KEY_BUF_NUM;
|
||||
static constexpr uint32_t M_MTE1_EVENT = EVENT_ID3;
|
||||
|
||||
static constexpr uint32_t MTE2_MTE1_EVENT = EVENT_ID2;
|
||||
static constexpr uint32_t MTE1_M_EVENT = EVENT_ID2;
|
||||
static constexpr uint32_t FIX_M_EVENT = EVENT_ID2;
|
||||
static constexpr uint32_t M_FIX_EVENT = EVENT_ID3;
|
||||
|
||||
static constexpr uint64_t M_BASIC_BLOCK = 256;
|
||||
static constexpr uint64_t D_BASIC_BLOCK = 128;
|
||||
static constexpr uint64_t S2_BASIC_BLOCK = 128;
|
||||
|
||||
static constexpr uint64_t M_BASIC_BLOCK_L0 = 128;
|
||||
static constexpr uint64_t D_BASIC_BLOCK_L0 = 128;
|
||||
static constexpr uint64_t S2_BASIC_BLOCK_L0 = 128;
|
||||
|
||||
static constexpr uint64_t FP16_BLOCK_CUBE = 16;
|
||||
static constexpr FixpipeConfig QLI_CFG_ROW_MAJOR_UB = {CO2Layout::ROW_MAJOR, true};
|
||||
|
||||
static constexpr uint64_t QUERY_BUFFER_OFFSET = M_BASIC_BLOCK * D_BASIC_BLOCK;
|
||||
static constexpr uint64_t KEY_BUFFER_OFFSET = S2_BASIC_BLOCK * D_BASIC_BLOCK;
|
||||
static constexpr uint64_t L0AB_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0;
|
||||
static constexpr uint64_t L0C_BUFFER_OFFSET = M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0;
|
||||
|
||||
protected:
|
||||
__aicore__ inline void Fixp(uint64_t s1gGmOffset, uint64_t s2GmOffset, uint64_t s1gL0RealSize,
|
||||
uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void ComputeL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void LoadKeyToL0b(uint64_t s2L0Offset, uint64_t s2L1RealSize, uint64_t s2L0RealSize,
|
||||
const LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void LoadQueryToL0a(uint64_t s1gL1Offset, uint64_t s1gL0Offset, uint64_t s1gL1RealSize,
|
||||
uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gL1Offset, const LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo);
|
||||
__aicore__ inline void KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset, const LICommon::RunInfo &runInfo);
|
||||
GlobalTensor<int32_t> blkTableGm_;
|
||||
GlobalTensor<K_T> keyGm_;
|
||||
GlobalTensor<Q_T> queryGm_;
|
||||
|
||||
TBuf<TPosition::A1> bufQL1_;
|
||||
LocalTensor<Q_T> queryL1_;
|
||||
TBuf<TPosition::B1> bufKeyL1_;
|
||||
LocalTensor<K_T> keyL1_;
|
||||
|
||||
TBuf<TPosition::A2> bufQL0_;
|
||||
LocalTensor<Q_T> queryL0_;
|
||||
TBuf<TPosition::B2> bufKeyL0_;
|
||||
LocalTensor<K_T> keyL0_;
|
||||
|
||||
TBuf<TPosition::CO1> bufL0C_;
|
||||
LocalTensor<float> cL0_;
|
||||
|
||||
TBuf<TPosition::VECCALC> bufUB_;
|
||||
LocalTensor<float> mm1ResUB_;
|
||||
|
||||
uint64_t keyL1BufIdx_ = 0;
|
||||
uint64_t queryL1Mte2BufIdx_ = 0;
|
||||
uint64_t queryL1Mte1BufIdx_ = 0;
|
||||
uint64_t l0BufIdx_ = 0;
|
||||
uint64_t kl0BufIdx_ = 0;
|
||||
|
||||
ConstInfo constInfo_;
|
||||
|
||||
private:
|
||||
static constexpr bool PAGE_ATTENTION = LIT::pageAttention;
|
||||
};
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::InitParams(const ConstInfo &constInfo)
|
||||
{
|
||||
constInfo_ = constInfo;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::InitBuffers(TPipe *pipe)
|
||||
{
|
||||
pipe->InitBuffer(bufUB_, 2 *CeilDiv(constInfo_.mBaseSize, 2) *
|
||||
constInfo_.s2BaseSize *
|
||||
sizeof(float)); // 大小：2(开dB) * 2 * 64 * 128 * 4 = 128KB
|
||||
mm1ResUB_ = bufUB_.Get<float>();
|
||||
pipe->InitBuffer(bufQL1_, QUERY_BUF_NUM * M_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(Q_T));
|
||||
queryL1_ = bufQL1_.Get<Q_T>();
|
||||
pipe->InitBuffer(bufKeyL1_, KEY_BUF_NUM * S2_BASIC_BLOCK * D_BASIC_BLOCK * sizeof(K_T));
|
||||
keyL1_ = bufKeyL1_.Get<K_T>();
|
||||
|
||||
pipe->InitBuffer(bufQL0_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * D_BASIC_BLOCK_L0 * sizeof(Q_T));
|
||||
queryL0_ = bufQL0_.Get<Q_T>();
|
||||
pipe->InitBuffer(bufKeyL0_, L0_BUF_NUM * D_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(K_T));
|
||||
keyL0_ = bufKeyL0_.Get<K_T>();
|
||||
|
||||
pipe->InitBuffer(bufL0C_, L0_BUF_NUM * M_BASIC_BLOCK_L0 * S2_BASIC_BLOCK_L0 * sizeof(float));
|
||||
cL0_ = bufL0C_.Get<float>();
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void
|
||||
LightningIndexerServiceCube<LIT>::InitMm1GlobalTensor(const GlobalTensor<int32_t> &blkTableGm,
|
||||
const GlobalTensor<K_T> &keyGm, const GlobalTensor<Q_T> &queryGm)
|
||||
{
|
||||
blkTableGm_ = blkTableGm;
|
||||
keyGm_ = keyGm;
|
||||
queryGm_ = queryGm;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::ComputeMm1(const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
CrossCoreWaitFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_FIX>(
|
||||
LICommon::ConstInfo::CROSS_VC_EVENT + runInfo.loop % 2);
|
||||
CrossCoreWaitFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_FIX>(
|
||||
LICommon::ConstInfo::CROSS_VC_EVENT +
|
||||
runInfo.loop % 2 +
|
||||
LICommon::ConstInfo::AIV0_AIV1_OFFSET);
|
||||
uint64_t s2GmBaseOffset = runInfo.s2Idx * constInfo_.s2BaseSize;
|
||||
uint64_t s1gProcessSize = runInfo.actMBaseSize;
|
||||
uint64_t s2ProcessSize = runInfo.actualSingleProcessSInnerSize;
|
||||
for (uint64_t s2GmOffset = 0; s2GmOffset < s2ProcessSize; s2GmOffset += S2_BASIC_BLOCK) {
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM);
|
||||
uint64_t s2L1RealSize =
|
||||
s2GmOffset + S2_BASIC_BLOCK > s2ProcessSize ? s2ProcessSize - s2GmOffset : S2_BASIC_BLOCK;
|
||||
if (PAGE_ATTENTION) {
|
||||
KeyNd2NzForPA(s2L1RealSize, s2GmBaseOffset + s2GmOffset, runInfo);
|
||||
}else {
|
||||
KeyNd2Nz(s2L1RealSize, s2GmOffset, runInfo);
|
||||
}
|
||||
|
||||
SetFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
|
||||
// s1gProcessSize当前必定不会超过2倍的s1g basic block
|
||||
for (uint64_t s1gGmOffset = 0; s1gGmOffset < s1gProcessSize; s1gGmOffset += constInfo_.mBaseSize) {
|
||||
uint64_t s1gL1RealSize =
|
||||
s1gGmOffset + constInfo_.mBaseSize > s1gProcessSize
|
||||
? s1gProcessSize - s1gGmOffset
|
||||
: constInfo_.mBaseSize;
|
||||
uint64_t s1gL1SizeAlign2G = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize);
|
||||
if (runInfo.isFirstS2InnerLoop && s2GmOffset == 0) {
|
||||
queryL1Mte2BufIdx_++;
|
||||
queryL1Mte1BufIdx_ = queryL1Mte2BufIdx_;
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + queryL1Mte2BufIdx_ % QUERY_BUF_NUM);
|
||||
QueryNd2Nz(s1gL1RealSize, s1gGmOffset, runInfo);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(MTE2_MTE1_EVENT);
|
||||
} else {
|
||||
queryL1Mte1BufIdx_ =
|
||||
queryL1Mte2BufIdx_ - (CeilDiv(s1gProcessSize, constInfo_.mBaseSize) - 1 - (s1gGmOffset > 0));
|
||||
}
|
||||
for (uint64_t s2L1Offset = 0; s2L1Offset < s2L1RealSize; s2L1Offset += S2_BASIC_BLOCK_L0) {
|
||||
uint64_t s2L0RealSize =
|
||||
s2L1Offset + S2_BASIC_BLOCK_L0 > s2L1RealSize ? s2L1RealSize - s2L1Offset : S2_BASIC_BLOCK_L0;
|
||||
|
||||
uint64_t l0Stride = constInfo_.mBaseSize;
|
||||
if (constInfo_.splitMFlag) {
|
||||
l0Stride /= 2;
|
||||
}
|
||||
|
||||
for (uint64_t s1gL1Offset = 0; s1gL1Offset < s1gL1SizeAlign2G; s1gL1Offset += l0Stride) {
|
||||
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM);
|
||||
uint64_t s1gL0RealSize =
|
||||
s1gL1Offset + constInfo_.mBaseSize > s1gL1SizeAlign2G
|
||||
? s1gL1SizeAlign2G - s1gL1Offset
|
||||
: constInfo_.mBaseSize;
|
||||
if (constInfo_.splitMFlag) {
|
||||
s1gL0RealSize = 128; // g=64, topK=2k时固定m=128
|
||||
}
|
||||
LoadQueryToL0a(s1gGmOffset, s1gL1Offset, s1gL1SizeAlign2G, s1gL0RealSize, runInfo);
|
||||
if (s1gL1Offset == 0) {
|
||||
LoadKeyToL0b(s2L1Offset, s2L1RealSize, s2L0RealSize, runInfo);
|
||||
}
|
||||
|
||||
SetFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
|
||||
WaitFlag<HardEvent::MTE1_M>(MTE1_M_EVENT);
|
||||
|
||||
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM);
|
||||
ComputeL0c(s1gL0RealSize, s2L0RealSize, runInfo);
|
||||
|
||||
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + l0BufIdx_ % L0_BUF_NUM);
|
||||
|
||||
bool lastIter = s1gL1Offset + l0Stride >= s1gL1SizeAlign2G;
|
||||
if (lastIter) {
|
||||
kl0BufIdx_++;
|
||||
}
|
||||
|
||||
Fixp(s1gGmOffset + s1gL1Offset, s2GmOffset + s2L1Offset, s1gL0RealSize, s2L0RealSize, runInfo);
|
||||
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + l0BufIdx_ % L0_BUF_NUM);
|
||||
l0BufIdx_++;
|
||||
}
|
||||
}
|
||||
if (s2GmOffset + S2_BASIC_BLOCK >= s2ProcessSize && runInfo.isLastS2InnerLoop) {
|
||||
SetFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + queryL1Mte1BufIdx_ % QUERY_BUF_NUM);
|
||||
}
|
||||
}
|
||||
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + keyL1BufIdx_ % KEY_BUF_NUM);
|
||||
keyL1BufIdx_++;
|
||||
}
|
||||
CrossCoreSetFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_FIX>(
|
||||
LICommon::ConstInfo::CROSS_CV_EVENT +
|
||||
runInfo.loop % 2);
|
||||
CrossCoreSetFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_FIX>(
|
||||
LICommon::ConstInfo::CROSS_CV_EVENT +
|
||||
runInfo.loop % 2 +
|
||||
LICommon::ConstInfo::AIV0_AIV1_OFFSET);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::KeyNd2Nz(uint64_t s2L1RealSize, uint64_t s2GmOffset,
|
||||
const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = s2L1RealSize; // 行数
|
||||
nd2nzPara.dValue = constInfo_.headDim;
|
||||
nd2nzPara.srcDValue = constInfo_.headDim;
|
||||
nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = 0;
|
||||
// 默认一块buf最多放两份
|
||||
DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET],
|
||||
keyGm_[runInfo.tensorKeyOffset + s2GmOffset * constInfo_.headDim], nd2nzPara);
|
||||
}
|
||||
|
||||
// blkNum, blkSize, N2, D
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::KeyNd2NzForPA(uint64_t s2L1RealSize, uint64_t s2GmOffset,
|
||||
const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
uint64_t s2L1Offset = 0;
|
||||
while (s2L1Offset < s2L1RealSize) {
|
||||
uint64_t s2BlkId = (s2L1Offset + s2GmOffset) / constInfo_.kCacheBlockSize;
|
||||
uint64_t s2BlkOffset = (s2L1Offset + s2GmOffset) % constInfo_.kCacheBlockSize;
|
||||
uint64_t keyGmOffset = blkTableGm_.GetValue(runInfo.bIdx * constInfo_.maxBlockNumPerBatch + s2BlkId) *
|
||||
constInfo_.kCacheBlockSize * constInfo_.kHeadNum * constInfo_.headDim +
|
||||
s2BlkOffset * constInfo_.headDim;
|
||||
|
||||
uint64_t s2Mte2Size = s2L1RealSize - s2L1Offset;
|
||||
s2Mte2Size = s2BlkOffset + s2Mte2Size >= constInfo_.kCacheBlockSize
|
||||
? constInfo_.kCacheBlockSize - s2BlkOffset
|
||||
: s2Mte2Size;
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = s2Mte2Size; // 行数
|
||||
nd2nzPara.dValue = constInfo_.headDim;
|
||||
nd2nzPara.srcDValue = constInfo_.headDim;
|
||||
nd2nzPara.dstNzC0Stride = CeilAlign(s2L1RealSize, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = 0;
|
||||
DataCopy(keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET + s2L1Offset * FP16_BLOCK_CUBE],
|
||||
keyGm_[keyGmOffset], nd2nzPara);
|
||||
|
||||
s2L1Offset += s2Mte2Size;
|
||||
}
|
||||
}
|
||||
|
||||
// batch, s1, n2, g, d
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::QueryNd2Nz(uint64_t s1gL1RealSize, uint64_t s1gGmOffset,
|
||||
const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
uint64_t dstNzC0Stride = CeilAlign(s1gL1RealSize, 2 * constInfo_.gSize);
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = s1gL1RealSize; // 行数
|
||||
nd2nzPara.dValue = constInfo_.headDim;
|
||||
nd2nzPara.srcDValue = constInfo_.headDim;
|
||||
nd2nzPara.dstNzC0Stride = CeilAlign(dstNzC0Stride, (uint64_t)BLOCK_CUBE); // 对齐到16 单位block
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = 0;
|
||||
// 默认一块buf最多放两份
|
||||
DataCopy(queryL1_[(queryL1Mte2BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET],
|
||||
queryGm_[runInfo.tensorQueryOffset + s1gGmOffset * constInfo_.headDim], nd2nzPara);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::LoadQueryToL0a(uint64_t s1gGmOffset,
|
||||
uint64_t s1gL1Offset, uint64_t s1gL1RealSize,
|
||||
uint64_t s1gL0RealSize, const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
LoadData2DParamsV2 loadData2DParamsV2;
|
||||
if (constInfo_.splitMFlag && runInfo.actMBaseSize > 128) { // 非尾块，切M
|
||||
uint64_t dstOffset = 0;
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = CeilDiv(64, BLOCK_CUBE);
|
||||
loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP16_BLOCK_CUBE);
|
||||
loadData2DParamsV2.srcStride = CeilDiv(s1gL1RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.dstStride = CeilDiv(s1gL0RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
for (int i = 0; i < 2; i++) {
|
||||
loadData2DParamsV2.mStartPosition = CeilDiv((s1gL1Offset / 2) + i * 128, BLOCK_CUBE);
|
||||
dstOffset = i * 64 * 16;
|
||||
|
||||
LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET + dstOffset],
|
||||
queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], loadData2DParamsV2);
|
||||
}
|
||||
} else {
|
||||
loadData2DParamsV2.mStartPosition = CeilDiv(s1gL1Offset, BLOCK_CUBE);
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = CeilDiv(s1gL0RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP16_BLOCK_CUBE);
|
||||
loadData2DParamsV2.srcStride = CeilDiv(s1gL1RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.dstStride = CeilDiv(s1gL0RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
|
||||
LoadData(queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET],
|
||||
queryL1_[(queryL1Mte1BufIdx_ % QUERY_BUF_NUM) * QUERY_BUFFER_OFFSET], loadData2DParamsV2);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::LoadKeyToL0b(uint64_t s2L1Offset,
|
||||
uint64_t s2L1RealSize, uint64_t s2L0RealSize,
|
||||
const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
LoadData2DParamsV2 loadData2DParamsV2;
|
||||
loadData2DParamsV2.mStartPosition = CeilDiv(s2L1Offset, BLOCK_CUBE);
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = CeilDiv(s2L0RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.kStep = CeilDiv(constInfo_.headDim, FP16_BLOCK_CUBE);
|
||||
loadData2DParamsV2.srcStride = CeilDiv(s2L1RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.dstStride = CeilDiv(s2L0RealSize, BLOCK_CUBE);
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
|
||||
LoadData(keyL0_[(kl0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET],
|
||||
keyL1_[(keyL1BufIdx_ % KEY_BUF_NUM) * KEY_BUFFER_OFFSET], loadData2DParamsV2);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::ComputeL0c(uint64_t s1gL0RealSize, uint64_t s2L0RealSize,
|
||||
const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
MmadParams mmadParams;
|
||||
mmadParams.m = CeilAlign(s1gL0RealSize, BLOCK_CUBE);
|
||||
mmadParams.n = s2L0RealSize;
|
||||
mmadParams.k = constInfo_.headDim;
|
||||
mmadParams.cmatrixInitVal = true;
|
||||
mmadParams.cmatrixSource = false;
|
||||
Mmad(cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], queryL0_[(l0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET],
|
||||
keyL0_[(kl0BufIdx_ % L0_BUF_NUM) * L0AB_BUFFER_OFFSET], mmadParams);
|
||||
if ((mmadParams.m / 16) * (mmadParams.n / 16) < 10) {
|
||||
PipeBarrier<PIPE_M>();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::Fixp(uint64_t s1gGmOffset,
|
||||
uint64_t s2GmOffset, uint64_t s1gL0RealSize,
|
||||
uint64_t s2L0RealSize, const LICommon::RunInfo &runInfo)
|
||||
{
|
||||
SetFlag<HardEvent::M_FIX>(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM);
|
||||
WaitFlag<HardEvent::M_FIX>(M_FIX_EVENT + l0BufIdx_ % L0_BUF_NUM);
|
||||
|
||||
static_assert(S2_BASIC_BLOCK == S2_BASIC_BLOCK_L0 && S2_BASIC_BLOCK_L0 == 128);
|
||||
// s1gL0RealSize：2*gSize(128)对齐, 最大256
|
||||
// s2L0RealSize <= S2_BASIC_BLOCK_L0, 未约束
|
||||
uint32_t nSize = (s2L0RealSize + 7) >> 3 << 3; // 32B对齐
|
||||
uint32_t mSize = (s1gL0RealSize + 1) >> 1 << 1;
|
||||
FixpipeParamsC310<CO2Layout::ROW_MAJOR> fixpipeParams;
|
||||
// 固定参数
|
||||
fixpipeParams.mSize = mSize;
|
||||
fixpipeParams.srcStride = mSize; // 已16对齐
|
||||
fixpipeParams.dstStride = UB_BANK_DEPTH_STRIDE / sizeof(float); // 落到同一个bank
|
||||
fixpipeParams.dualDstCtl = 1; // 双目标模式，按M维度拆分， M / 2 * N写入每个UB，M必须为2的倍数
|
||||
|
||||
uint64_t dstOffset = 0;
|
||||
if (constInfo_.splitMFlag && runInfo.actMBaseSize > 128) { // 非尾块，切M
|
||||
dstOffset = s1gGmOffset * 64;
|
||||
}
|
||||
|
||||
// nSize已保证N方向32B对齐
|
||||
if (nSize <= (256 / sizeof(float))) {
|
||||
// N方向小于一个bank(256B), 只需搬一个ND块, 且不用补齐
|
||||
fixpipeParams.nSize = nSize;
|
||||
fixpipeParams.params.ndNum = 1;
|
||||
fixpipeParams.params.srcNdStride = 0;
|
||||
fixpipeParams.params.dstNdStride = 0;
|
||||
} else {
|
||||
// N方向在(256B, 512B]范围， 直接按512B搬, 注意此时不能开unitflag
|
||||
fixpipeParams.nSize = S2_BASIC_BLOCK_L0 / 2; // 分2个ND搬, S2_BASIC_BLOCK_L0不为128会有问题
|
||||
fixpipeParams.params.ndNum = 2;
|
||||
fixpipeParams.params.srcNdStride = ((fixpipeParams.mSize + 15) / 16) * fixpipeParams.nSize;
|
||||
fixpipeParams.params.dstNdStride = constInfo_.s2BaseSize * constInfo_.mBaseSize / 2;
|
||||
}
|
||||
Fixpipe<float, float, QLI_CFG_ROW_MAJOR_UB>(mm1ResUB_[(runInfo.loop % 2) * constInfo_.s2BaseSize / 2 + dstOffset],
|
||||
cL0_[(l0BufIdx_ % L0_BUF_NUM) * L0C_BUFFER_OFFSET], fixpipeParams);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::AllocEventID()
|
||||
{
|
||||
SetMMLayoutTransform(true);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 1);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 2);
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 1);
|
||||
|
||||
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 0);
|
||||
SetFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 1);
|
||||
|
||||
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + 0);
|
||||
SetFlag<HardEvent::FIX_M>(FIX_M_EVENT + 1);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceCube<LIT>::FreeEventID()
|
||||
{
|
||||
SetMMLayoutTransform(false);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 1);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(KEY_MTE1_MTE2_EVENT + 2);
|
||||
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(QUERY_MTE1_MTE2_EVENT + 1);
|
||||
|
||||
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 0);
|
||||
WaitFlag<HardEvent::M_MTE1>(M_MTE1_EVENT + 1);
|
||||
|
||||
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + 0);
|
||||
WaitFlag<HardEvent::FIX_M>(FIX_M_EVENT + 1);
|
||||
}
|
||||
} // namespace LIKernel
|
||||
#endif
|
||||
@@ -0,0 +1,575 @@
|
||||
/**
|
||||
* 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 lightning_indexer_service_vector.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef LIGHTNING_INDEXER_SERVICE_VECTOR_H
|
||||
#define LIGHTNING_INDEXER_SERVICE_VECTOR_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "kernel_operator_list_tensor_intf.h"
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "lib/matrix/matmul/tiling.h"
|
||||
#include "../lightning_indexer_common.h"
|
||||
#include "../arch35/vf/lightning_indexer_vector1.h"
|
||||
#include "../arch35/vf/lightning_indexer_topk.h"
|
||||
|
||||
namespace LIKernel {
|
||||
using namespace LICommon;
|
||||
constexpr uint32_t TRUNK_LEN_16K = 16384;
|
||||
constexpr uint32_t TRUNK_LEN_8K = 8192;
|
||||
constexpr uint32_t TOPK_LEN_4K = 4096;
|
||||
|
||||
template<typename Q_T, typename W_T = void>
|
||||
struct LightningIndexerTypeTraits {
|
||||
using weightsType = Q_T; // 默认：weightsType绑定Q_T
|
||||
};
|
||||
|
||||
template<typename Q_T>
|
||||
struct LightningIndexerTypeTraits<Q_T, float> {
|
||||
using weightsType = float; // W_T=float时，强制weightsType为float
|
||||
};
|
||||
template <typename LIT>
|
||||
class LightningIndexerServiceVector {
|
||||
public:
|
||||
// =================================类型定义区=================================
|
||||
static constexpr LI_LAYOUT LAYOUT_T = LIT::layout;
|
||||
static constexpr LI_LAYOUT K_LAYOUT_T = LIT::keyLayout;
|
||||
static constexpr bool PAGE_ATTENTION = LIT::pageAttention;
|
||||
static constexpr bool DT_W_FLAG = LIT::weightsTypeFlag;
|
||||
using Q_T = typename LIT::queryType;
|
||||
using K_T = typename LIT::keyType;
|
||||
using W_T = typename LightningIndexerTypeTraits<Q_T,
|
||||
typename std::conditional<DT_W_FLAG, float, void>::type>::weightsType;
|
||||
|
||||
__aicore__ inline LightningIndexerServiceVector(){};
|
||||
__aicore__ inline void ProcessVec1(const LICommon::RunInfo &info);
|
||||
__aicore__ inline void ProcessTopK(const LICommon::RunInfo &info);
|
||||
__aicore__ inline void InitBuffers(TPipe *pipe);
|
||||
__aicore__ inline void InitParams(const struct LICommon::ConstInfo &constInfo,
|
||||
const LITilingData *__restrict tilingData);
|
||||
__aicore__ inline void InitVecWorkspaceTensor(GlobalTensor<uint16_t> scoreGm);
|
||||
__aicore__ inline void InitVecInputTensor(GlobalTensor<W_T> weightsGm, GlobalTensor<int32_t> indiceOutGm,
|
||||
GlobalTensor<K_T> valueOutGm, GlobalTensor<int32_t> blockTableGm);
|
||||
__aicore__ inline void CleanInvalidOutput(int64_t invalidS1offset);
|
||||
__aicore__ inline void AllocEventID();
|
||||
__aicore__ inline void FreeEventID();
|
||||
|
||||
protected:
|
||||
GlobalTensor<uint16_t> scoreGm;
|
||||
GlobalTensor<W_T> weightsGm;
|
||||
GlobalTensor<int32_t> indiceOutGm;
|
||||
GlobalTensor<K_T> valueOutGm;
|
||||
GlobalTensor<int32_t> blockTableGm;
|
||||
// =================================常量区=================================
|
||||
static constexpr uint32_t VEC1_V_MTE2_EVENT = EVENT_ID0;
|
||||
static constexpr uint32_t VEC1_MTE2_V_EVENT = EVENT_ID1;
|
||||
static constexpr uint32_t VEC1_V_MTE3_EVENT = EVENT_ID2;
|
||||
static constexpr uint32_t VEC1_MTE3_V_EVENT = EVENT_ID3;
|
||||
|
||||
static constexpr uint32_t TOPK_V_MTE2_EVENT = EVENT_ID4;
|
||||
static constexpr uint32_t TOPK_MTE2_V_EVENT = EVENT_ID5;
|
||||
static constexpr uint32_t TOPK_V_MTE3_EVENT = EVENT_ID6;
|
||||
static constexpr uint32_t TOPK_MTE3_V_EVENT = EVENT_ID7;
|
||||
|
||||
static constexpr uint32_t MTE3_MTE2_EVENT = EVENT_ID0;
|
||||
static constexpr uint32_t V_MTE2_EVENT = EVENT_ID7;
|
||||
static constexpr uint32_t V_MTE2_EVENT1 = EVENT_ID2;
|
||||
static constexpr uint32_t V_MTE2_EVENT2 = EVENT_ID3;
|
||||
static constexpr uint32_t V_MTE2_EVENT3 = EVENT_ID5;
|
||||
|
||||
private:
|
||||
// ================================Local Buffer区====================================
|
||||
|
||||
// tmp buff for vector
|
||||
TBuf<TPosition::VECCALC> resMm1Buf_;
|
||||
LocalTensor<float> resMm1UB_;
|
||||
// tmp buff for weight
|
||||
TBuf<TPosition::VECCALC> weightBuf_;
|
||||
LocalTensor<W_T> weightUB_;
|
||||
// tmp buff for weight cast float
|
||||
TBuf<TPosition::VECCALC> weightFloatBuf_;
|
||||
LocalTensor<float> weightFloatUB_;
|
||||
|
||||
// tmp buff for out
|
||||
TBuf<TPosition::VECCALC> outBuf_;
|
||||
LocalTensor<uint16_t> vec1OutUB_;
|
||||
|
||||
// tmp buff for returnValue K_T
|
||||
TBuf<TPosition::VECCALC> valueOutBuf_;
|
||||
LocalTensor<K_T> valueOutLocal_;
|
||||
|
||||
// tmp buff for topk
|
||||
TBuf<TPosition::VECCALC> mrgValueBuf_;
|
||||
LocalTensor<uint16_t> mrgValueLocal_;
|
||||
|
||||
TBuf<TPosition::VECCALC> indicesOutBuf_;
|
||||
LocalTensor<uint32_t> indicesOutLocal_;
|
||||
|
||||
TBuf<TPosition::VECCALC> scoreOutBuf_;
|
||||
LocalTensor<uint16_t> scoreOutLocal_;
|
||||
|
||||
TBuf<TPosition::VECCALC> topkSharedTmpBuf_;
|
||||
LocalTensor<uint32_t> topkSharedTmpLocal_;
|
||||
|
||||
int32_t blockId_ = -1;
|
||||
// para for vector
|
||||
int32_t groupInner_ = 0;
|
||||
int32_t globalTopkNum_ = 0;
|
||||
int64_t blockS2StartIdx_ = 0;
|
||||
int32_t gSize_ = 0;
|
||||
int32_t kSeqSize_ = 0;
|
||||
int32_t kHeadNum_ = 0;
|
||||
int32_t qHeadNum_ = 0;
|
||||
int32_t s1BaseSize_ = 0;
|
||||
int32_t s2BaseSize_ = 0;
|
||||
int32_t kCacheBlockSize_ = 0;
|
||||
int32_t maxBlockNumPerBatch_ = 0;
|
||||
uint32_t topkCount_ = 0;
|
||||
uint32_t topkCountAlign256_ = 0; // topkCount对齐到256(直方图需要)，支持topk泛化
|
||||
uint32_t trunkLen_ = 0;
|
||||
bool returnValueFlag = false;
|
||||
|
||||
struct LICommon::ConstInfo constInfo_;
|
||||
topk::LITopk<uint16_t> topkOp_;
|
||||
};
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::InitBuffers(TPipe *pipe)
|
||||
{
|
||||
pipe->InitBuffer(resMm1Buf_, 2 * CeilDiv(constInfo_.mBaseSize, 2) * s2BaseSize_ * sizeof(float));
|
||||
resMm1UB_ = resMm1Buf_.Get<float>();
|
||||
|
||||
pipe->InitBuffer(weightBuf_, 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE);
|
||||
weightUB_ = weightBuf_.Get<W_T>();
|
||||
pipe->InitBuffer(weightFloatBuf_, 2 * CeilDiv(s1BaseSize_, 2) * UB_BANK_DEPTH_STRIDE);
|
||||
weightFloatUB_ = weightFloatBuf_.Get<float>();
|
||||
pipe->InitBuffer(outBuf_,
|
||||
2 * CeilDiv(s1BaseSize_, 2) * s2BaseSize_ * sizeof(uint16_t)); // 大小：2(开dB) * 2 * 128 * 4 = 2KB
|
||||
vec1OutUB_ = outBuf_.Get<uint16_t>(); // out
|
||||
|
||||
// Topk
|
||||
pipe->InitBuffer(mrgValueBuf_,
|
||||
(topkCountAlign256_ + trunkLen_) * sizeof(uint16_t));
|
||||
mrgValueLocal_ = mrgValueBuf_.Get<uint16_t>();
|
||||
// returnvalue
|
||||
if (topkCount_ <= 2048) {
|
||||
pipe->InitBuffer(valueOutBuf_, topkCountAlign256_ * sizeof(K_T));
|
||||
valueOutLocal_ = valueOutBuf_.Get<K_T>();
|
||||
} else { // sparseCount > 2k时，复用return value相关UB
|
||||
valueOutLocal_ = mrgValueBuf_.Get<K_T>(); // returnValue float
|
||||
}
|
||||
|
||||
// 大小：(topkCountAlign256_ + 64) * 4 64:duplicate刷-1需要额外空间
|
||||
pipe->InitBuffer(indicesOutBuf_,
|
||||
(topkCountAlign256_ + 64) * sizeof(uint32_t));
|
||||
indicesOutLocal_ = indicesOutBuf_.Get<uint32_t>();
|
||||
|
||||
pipe->InitBuffer(scoreOutBuf_, topkCountAlign256_ * sizeof(uint16_t));
|
||||
scoreOutLocal_ = scoreOutBuf_.Get<uint16_t>();
|
||||
|
||||
uint64_t topkSharedTmpSize = topkOp_.GetSharedTmpBufferSize();
|
||||
pipe->InitBuffer(topkSharedTmpBuf_, topkSharedTmpSize);
|
||||
topkSharedTmpLocal_ = topkSharedTmpBuf_.Get<uint32_t>();
|
||||
topkOp_.InitBuffers(topkSharedTmpLocal_);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::InitParams(const struct LICommon::ConstInfo &constInfo,
|
||||
const LITilingData *__restrict tilingData)
|
||||
{
|
||||
this->constInfo_ = constInfo;
|
||||
blockS2StartIdx_ = 0;
|
||||
gSize_ = constInfo.gSize;
|
||||
kSeqSize_ = constInfo.kSeqSize;
|
||||
// define N2 para
|
||||
kHeadNum_ = constInfo.kHeadNum;
|
||||
qHeadNum_ = constInfo.qHeadNum;
|
||||
// define MMBase para
|
||||
s1BaseSize_ = constInfo.s1BaseSize; // 4
|
||||
s2BaseSize_ = constInfo.s2BaseSize; // 128
|
||||
kCacheBlockSize_ = constInfo.kCacheBlockSize;
|
||||
maxBlockNumPerBatch_ = constInfo.maxBlockNumPerBatch;
|
||||
returnValueFlag = constInfo.returnValueFlag;
|
||||
blockId_ = GetBlockIdx();
|
||||
trunkLen_ = constInfo.sparseCount >= TOPK_LEN_4K ? TRUNK_LEN_8K : TRUNK_LEN_16K;
|
||||
topkCount_ = constInfo.sparseCount;
|
||||
topkOp_.Init(topkCount_, trunkLen_);
|
||||
topkCountAlign256_ = LICommon::Align(constInfo.sparseCount, (uint64_t)256); // topkCount对齐到256
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::InitVecInputTensor(GlobalTensor<W_T> weightsGm,
|
||||
GlobalTensor<int32_t> indiceOutGm,
|
||||
GlobalTensor<K_T> valueOutGm,
|
||||
GlobalTensor<int32_t> blockTableGm)
|
||||
{
|
||||
this->weightsGm = weightsGm;
|
||||
this->indiceOutGm = indiceOutGm;
|
||||
this->valueOutGm = valueOutGm;
|
||||
this->blockTableGm = blockTableGm;
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::InitVecWorkspaceTensor(GlobalTensor<uint16_t> scoreGm)
|
||||
{
|
||||
this->scoreGm = scoreGm; // resucesum*k
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::AllocEventID()
|
||||
{
|
||||
SetFlag<HardEvent::V_MTE2>(VEC1_V_MTE2_EVENT + 0);
|
||||
SetFlag<HardEvent::V_MTE2>(VEC1_V_MTE2_EVENT + 1);
|
||||
SetFlag<HardEvent::MTE3_V>(VEC1_MTE3_V_EVENT + 0);
|
||||
SetFlag<HardEvent::MTE3_V>(VEC1_MTE3_V_EVENT + 1);
|
||||
|
||||
SetFlag<HardEvent::V_MTE2>(TOPK_V_MTE2_EVENT);
|
||||
SetFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT1);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::FreeEventID()
|
||||
{
|
||||
WaitFlag<HardEvent::V_MTE2>(VEC1_V_MTE2_EVENT + 0);
|
||||
WaitFlag<HardEvent::V_MTE2>(VEC1_V_MTE2_EVENT + 1);
|
||||
WaitFlag<HardEvent::MTE3_V>(VEC1_MTE3_V_EVENT + 0);
|
||||
WaitFlag<HardEvent::MTE3_V>(VEC1_MTE3_V_EVENT + 1);
|
||||
|
||||
WaitFlag<HardEvent::V_MTE2>(TOPK_V_MTE2_EVENT);
|
||||
WaitFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT1);
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::CleanInvalidOutput(int64_t invalidS1Offset)
|
||||
{
|
||||
// init -1 and copy to output
|
||||
uint64_t dealSize = constInfo_.sparseCount;
|
||||
GlobalTensor<int32_t> indexOutput = indiceOutGm[invalidS1Offset];
|
||||
AscendC::InitGlobalMemory(indexOutput, dealSize, constInfo_.INVALID_IDX);
|
||||
if (returnValueFlag) {
|
||||
SetFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
WaitFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
Duplicate(valueOutLocal_.template ReinterpretCast<uint16_t>(), constInfo_.INVALID_VAL, constInfo_.sparseCount);
|
||||
|
||||
SetFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
|
||||
AscendC::DataCopyParams copyOutValueParams;
|
||||
copyOutValueParams.blockCount = 1;
|
||||
copyOutValueParams.blockLen = constInfo_.sparseCount * sizeof(K_T);
|
||||
copyOutValueParams.srcStride = 0;
|
||||
copyOutValueParams.dstStride = 0;
|
||||
AscendC::DataCopyPad(valueOutGm[invalidS1Offset], valueOutLocal_, copyOutValueParams);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::ProcessVec1(const LICommon::RunInfo &info)
|
||||
{
|
||||
auto pingpong = (info.loop % 2);
|
||||
auto s1BaseSizePerAIV = CeilDiv(s1BaseSize_, 2);
|
||||
int64_t curS1Idx = info.gS1Idx * s1BaseSize_;
|
||||
int64_t curS2Idx = info.s2Idx * s2BaseSize_;
|
||||
int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_;
|
||||
int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2);
|
||||
int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2;
|
||||
if (curAivS1ProcNum == 0) {
|
||||
CrossCoreWaitFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_V>(
|
||||
LICommon::ConstInfo::CROSS_CV_EVENT + pingpong
|
||||
); // V核等C核计算完mm1，mm1Res已搬运到UB
|
||||
CrossCoreSetFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_V>(
|
||||
LICommon::ConstInfo::CROSS_VC_EVENT + pingpong
|
||||
); // V核处理完，通知C核可以把mm1Res搬运到UB
|
||||
return;
|
||||
}
|
||||
WaitFlag<HardEvent::V_MTE2>(VEC1_V_MTE2_EVENT + pingpong);
|
||||
// weightsGm --> weightUB_
|
||||
int64_t weightGmOffset = info.tensorWeightsOffset + curAivS1Idx * kHeadNum_ * gSize_;
|
||||
DataCopyPadExtParams<W_T> padWeightsParams{false, 0, 0, 0};
|
||||
DataCopyExtParams wDataCopyExtParams;
|
||||
wDataCopyExtParams.blockCount = curAivS1ProcNum;
|
||||
wDataCopyExtParams.blockLen = gSize_ * sizeof(W_T);
|
||||
wDataCopyExtParams.srcStride = 0;
|
||||
wDataCopyExtParams.dstStride = (UB_BANK_DEPTH_STRIDE - wDataCopyExtParams.blockLen) / 32;
|
||||
DataCopyPad(weightUB_[pingpong * (UB_BANK_STRIDE / sizeof(W_T))],
|
||||
weightsGm[weightGmOffset], wDataCopyExtParams, padWeightsParams);
|
||||
|
||||
SetFlag<HardEvent::MTE2_V>(VEC1_MTE2_V_EVENT + pingpong);
|
||||
WaitFlag<HardEvent::MTE2_V>(VEC1_MTE2_V_EVENT + pingpong);
|
||||
WaitFlag<HardEvent::MTE3_V>(VEC1_MTE3_V_EVENT + pingpong);
|
||||
|
||||
// CV同步
|
||||
CrossCoreWaitFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_V>(
|
||||
LICommon::ConstInfo::CROSS_CV_EVENT + info.loop % 2
|
||||
); // V核等C核计算完mm1，mm1Res已搬运到UB
|
||||
|
||||
auto outBase = vec1OutUB_[pingpong * (UB_BANK_STRIDE / sizeof(uint16_t))];
|
||||
auto weightBase = weightUB_[pingpong * (UB_BANK_STRIDE / sizeof(W_T))];
|
||||
auto weightFloatBase = weightFloatUB_[pingpong * (UB_BANK_STRIDE / sizeof(float))];
|
||||
auto qkBase = resMm1UB_[pingpong * (UB_BANK_STRIDE / sizeof(float))];
|
||||
auto qkVLstride = (UB_BANK_DEPTH_STRIDE / sizeof(float)) / 2 * constInfo_.mBaseSize;
|
||||
|
||||
vector1::BatchMulWeightAndReduceSum(outBase, UB_BANK_DEPTH_STRIDE / sizeof(uint16_t),
|
||||
qkBase, qkVLstride, (uint32_t)(gSize_ * UB_BANK_DEPTH_STRIDE / sizeof(float)),
|
||||
weightBase, UB_BANK_DEPTH_STRIDE / sizeof(W_T), weightFloatBase,
|
||||
gSize_, curAivS1ProcNum);
|
||||
SetFlag<HardEvent::V_MTE2>(VEC1_V_MTE2_EVENT + pingpong);
|
||||
SetFlag<HardEvent::V_MTE3>(VEC1_V_MTE3_EVENT + pingpong);
|
||||
WaitFlag<HardEvent::V_MTE3>(VEC1_V_MTE3_EVENT + pingpong);
|
||||
// outUB_ ---> scoreGm
|
||||
int64_t vec1OutGmOffset = blockId_ % 2 == 0
|
||||
? curS2Idx
|
||||
: s1BaseSizePerAIV * LICommon::Align(
|
||||
(uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_
|
||||
) + curS2Idx;
|
||||
DataCopyExtParams copyOutParams;
|
||||
copyOutParams.blockCount = curAivS1ProcNum;
|
||||
copyOutParams.blockLen = s2BaseSize_ * sizeof(uint16_t);
|
||||
copyOutParams.srcStride = (UB_BANK_DEPTH_STRIDE - UB_BANK_STRIDE) / 32;
|
||||
copyOutParams.dstStride = (LICommon::Align(
|
||||
(uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_
|
||||
) - s2BaseSize_) * sizeof(uint16_t);
|
||||
DataCopyPad(scoreGm[vec1OutGmOffset], outBase, copyOutParams);
|
||||
SetFlag<HardEvent::MTE3_V>(VEC1_MTE3_V_EVENT + pingpong);
|
||||
CrossCoreSetFlag<LICommon::ConstInfo::QLI_SYNC_MODE4, PIPE_V>(
|
||||
LICommon::ConstInfo::CROSS_VC_EVENT + pingpong
|
||||
); // V核处理完，通知C核可以把mm1Res搬运到UB
|
||||
}
|
||||
|
||||
template <typename LIT>
|
||||
__aicore__ inline void LightningIndexerServiceVector<LIT>::ProcessTopK(const LICommon::RunInfo &info)
|
||||
{
|
||||
SetFlag<HardEvent::MTE3_MTE2>(MTE3_MTE2_EVENT);
|
||||
WaitFlag<HardEvent::MTE3_MTE2>(MTE3_MTE2_EVENT);
|
||||
|
||||
int64_t curS1Idx = info.gS1Idx * s1BaseSize_;
|
||||
int64_t curS2Idx = info.s2Idx * s2BaseSize_;
|
||||
int64_t curS1ProcNum = curS1Idx + s1BaseSize_ > info.actS1Size ? info.actS1Size % s1BaseSize_ : s1BaseSize_;
|
||||
int64_t curAivS1Idx = curS1Idx + (blockId_ % 2) * CeilDiv(curS1ProcNum, 2);
|
||||
int64_t curAivS1ProcNum = (blockId_ % 2 == 0) ? CeilDiv(curS1ProcNum, 2) : curS1ProcNum / 2;
|
||||
|
||||
AscendC::DataCopyExtParams copyInParams;
|
||||
copyInParams.blockCount = 1;
|
||||
copyInParams.srcStride = 0;
|
||||
copyInParams.dstStride = 0;
|
||||
copyInParams.rsv = 0;
|
||||
|
||||
AscendC::DataCopyParams copyOutParams;
|
||||
copyOutParams.blockCount = 1;
|
||||
copyOutParams.blockLen = topkCount_ * sizeof(uint32_t); // bytes
|
||||
copyOutParams.srcStride = 0;
|
||||
copyOutParams.dstStride = 0;
|
||||
|
||||
int32_t cuRealAcSeq = info.actS2Size;
|
||||
if (constInfo_.attenMaskFlag) {
|
||||
cuRealAcSeq = info.actS2SizeOrig - info.actS1Size + curAivS1Idx + 1;
|
||||
}
|
||||
|
||||
int32_t validS2Len = cuRealAcSeq;
|
||||
for (uint32_t i = 0; i < curAivS1ProcNum; i++) {
|
||||
uint32_t rowIdx = blockId_ % 2 * CeilDiv(curS1ProcNum, 2) + i;
|
||||
uint32_t vecOffset = blockId_ % 2 * CeilDiv(s1BaseSize_, 2) + i;
|
||||
|
||||
uint16_t zero = 0;
|
||||
int32_t neg = -1;
|
||||
if (constInfo_.attenMaskFlag) {
|
||||
validS2Len = (int32_t)i + cuRealAcSeq;
|
||||
}
|
||||
if (validS2Len <= 0) {
|
||||
WaitFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
Duplicate(indicesOutLocal_.ReinterpretCast<int32_t>(), neg, topkCount_);
|
||||
SetFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_],
|
||||
indicesOutLocal_.ReinterpretCast<int32_t>(),
|
||||
copyOutParams);
|
||||
SetFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
if (returnValueFlag) {
|
||||
WaitFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
Duplicate(valueOutLocal_.template ReinterpretCast<uint16_t>(), constInfo_.INVALID_VAL, topkCount_);
|
||||
|
||||
SetFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
|
||||
AscendC::DataCopyParams copyOutValueParams;
|
||||
copyOutValueParams.blockCount = 1;
|
||||
copyOutValueParams.blockLen = topkCount_ * sizeof(K_T);
|
||||
copyOutValueParams.srcStride = 0;
|
||||
copyOutValueParams.dstStride = 0;
|
||||
AscendC::DataCopyPad(
|
||||
valueOutGm[info.valueOutOffset + (curS1Idx + rowIdx) * topkCount_],
|
||||
valueOutLocal_,
|
||||
copyOutValueParams);
|
||||
SetFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
}
|
||||
continue;
|
||||
}
|
||||
|
||||
WaitFlag<HardEvent::V_MTE2>(TOPK_V_MTE2_EVENT);
|
||||
WaitFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
|
||||
AscendC::DataCopyPadExtParams<uint16_t> padParams{true, 0, 0, 0};
|
||||
if (validS2Len >= topkCount_) {
|
||||
uint32_t s2LoopNum = (validS2Len + trunkLen_ - 1) / trunkLen_;
|
||||
if (s2LoopNum == 1) {
|
||||
uint32_t validS2LenAlign = LICommon::Align(validS2Len, (int32_t)256);
|
||||
Duplicate(mrgValueLocal_[validS2Len / 256 * 256], zero, validS2LenAlign - validS2Len / 256 * 256);
|
||||
SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT);
|
||||
copyInParams.blockLen = validS2Len * sizeof(uint16_t); // byte
|
||||
AscendC::DataCopyPadExtParams<uint16_t> padParams{true, 0, 0, 0};
|
||||
AscendC::DataCopyPad(
|
||||
mrgValueLocal_,
|
||||
scoreGm[vecOffset * LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)],
|
||||
copyInParams, padParams);
|
||||
SetFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
WaitFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
topkOp_(mrgValueLocal_, indicesOutLocal_, scoreOutLocal_, validS2LenAlign, 0, 1, returnValueFlag);
|
||||
} else {
|
||||
for (uint32_t loopIdx = 0; loopIdx < s2LoopNum; loopIdx++) {
|
||||
if (loopIdx == 0) {
|
||||
copyInParams.blockLen = trunkLen_ * sizeof(uint16_t); // byte
|
||||
AscendC::DataCopyPad(
|
||||
mrgValueLocal_,
|
||||
scoreGm[vecOffset * LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)],
|
||||
copyInParams, padParams);
|
||||
SetFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
WaitFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
topkOp_(mrgValueLocal_, indicesOutLocal_,
|
||||
scoreOutLocal_, trunkLen_, loopIdx,
|
||||
s2LoopNum, returnValueFlag);
|
||||
continue;
|
||||
}
|
||||
SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT2);
|
||||
WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT2);
|
||||
uint32_t validTrunkLen = (loopIdx * trunkLen_ + trunkLen_) > validS2Len
|
||||
? validS2Len % trunkLen_
|
||||
:trunkLen_;
|
||||
uint32_t offset = vecOffset *
|
||||
LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_) +
|
||||
loopIdx * trunkLen_;
|
||||
AscendC::DataCopy(mrgValueLocal_, scoreOutLocal_, topkCountAlign256_);
|
||||
// topk如果没有对齐到256，则把topkCountAlign256_ - topkCount_部分刷0
|
||||
if (topkCountAlign256_ != topkCount_) {
|
||||
uint64_t mask[1];
|
||||
mask[0] = ~0;
|
||||
mask[0] = mask[0] << (topkCount_ % 64);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// 把topkCount_对齐到64刷0，此处由于duplicate的限制mask[0]刷64个数
|
||||
Duplicate(mrgValueLocal_[topkCount_ / 64 * 64], zero, mask, 1, 1, 0);
|
||||
PipeBarrier<PIPE_V>();
|
||||
// 把topk剩余对齐到256的部分刷0
|
||||
Duplicate(mrgValueLocal_[topkCount_ / 64 * 64 + 64], zero,
|
||||
topkCountAlign256_ - (topkCount_ / 64 * 64 + 64));
|
||||
SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT3);
|
||||
WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT3);
|
||||
}
|
||||
copyInParams.blockLen = validTrunkLen * sizeof(uint16_t); // byte
|
||||
// TOPK 直方图一次必须计算256，输入处理数据需要和256对齐
|
||||
if ((topkCountAlign256_ + validTrunkLen) % 256 != 0) {
|
||||
Duplicate(mrgValueLocal_[topkCountAlign256_ + validTrunkLen / 256 * 256],
|
||||
zero, LICommon::Align(validTrunkLen,
|
||||
(uint32_t)256) - validTrunkLen / 256 * 256);
|
||||
SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT);
|
||||
}
|
||||
WaitFlag<HardEvent::V_MTE2>(V_MTE2_EVENT1);
|
||||
AscendC::DataCopyPad(mrgValueLocal_[topkCountAlign256_], scoreGm[offset], copyInParams, padParams);
|
||||
SetFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
WaitFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
topkOp_(mrgValueLocal_, indicesOutLocal_,
|
||||
scoreOutLocal_,
|
||||
LICommon::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256),
|
||||
loopIdx, s2LoopNum, returnValueFlag);
|
||||
SetFlag<HardEvent::V_MTE2>(V_MTE2_EVENT1);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
AscendC::CreateVecIndex(indicesOutLocal_.ReinterpretCast<int32_t>(), (int32_t)zero, validS2Len);
|
||||
if (returnValueFlag) {
|
||||
copyInParams.blockLen = LICommon::Align(validS2Len, (int32_t)32) * sizeof(uint16_t);
|
||||
AscendC::DataCopyPad(scoreOutLocal_,
|
||||
scoreGm[vecOffset * LICommon::Align((uint64_t)constInfo_.kSeqSize, (uint64_t)s2BaseSize_)],
|
||||
copyInParams, padParams);
|
||||
SetFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
WaitFlag<HardEvent::MTE2_V>(TOPK_MTE2_V_EVENT);
|
||||
}
|
||||
}
|
||||
|
||||
if (validS2Len < topkCount_) {
|
||||
uint64_t mask[1];
|
||||
mask[0] = ~0;
|
||||
mask[0] = mask[0] << (validS2Len % 8);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Duplicate(indicesOutLocal_.ReinterpretCast<int32_t>()[validS2Len / 8 * 8], neg, mask, 1, 1, 0);
|
||||
}
|
||||
|
||||
if (validS2Len / 8 * 8 + 64 < topkCount_) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Duplicate(indicesOutLocal_.ReinterpretCast<int32_t>()[validS2Len / 8 * 8 + 64],
|
||||
neg, topkCount_ - (validS2Len / 8 * 8 + 64));
|
||||
}
|
||||
|
||||
SetFlag<HardEvent::V_MTE2>(TOPK_V_MTE2_EVENT);
|
||||
SetFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
AscendC::DataCopyPad(indiceOutGm[info.indiceOutOffset + (curS1Idx + rowIdx) * topkCount_],
|
||||
indicesOutLocal_.ReinterpretCast<int32_t>(), copyOutParams);
|
||||
|
||||
|
||||
// 是否返回Value值
|
||||
if (returnValueFlag) {
|
||||
WaitFlag<HardEvent::V_MTE2>(TOPK_V_MTE2_EVENT);
|
||||
// uint16_t -> bfloat16
|
||||
if (std::is_same_v<K_T, bfloat16_t>) {
|
||||
vector1::UIntToFloatReturnValue(valueOutLocal_.template ReinterpretCast<bfloat16_t>(),
|
||||
scoreOutLocal_, topkCountAlign256_);
|
||||
} else {
|
||||
vector1::UIntToFloatReturnValue(valueOutLocal_.template ReinterpretCast<half>(),
|
||||
scoreOutLocal_, topkCountAlign256_);
|
||||
}
|
||||
|
||||
if (validS2Len < topkCount_) {
|
||||
uint64_t mask[1];
|
||||
mask[0] = ~0;
|
||||
mask[0] = mask[0] << (validS2Len % 16);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Duplicate(valueOutLocal_.template ReinterpretCast<uint16_t>()[validS2Len / 16 * 16],
|
||||
constInfo_.INVALID_VAL, mask, 1, 1, 0);
|
||||
}
|
||||
if (validS2Len / 16 * 16 + 64 < topkCount_) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
Duplicate(valueOutLocal_.template ReinterpretCast<uint16_t>()[validS2Len / 16 * 16 + 64],
|
||||
constInfo_.INVALID_VAL, topkCount_ - (validS2Len / 16 * 16 + 64));
|
||||
}
|
||||
SetFlag<HardEvent::V_MTE2>(TOPK_V_MTE2_EVENT);
|
||||
SetFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
WaitFlag<HardEvent::V_MTE3>(TOPK_V_MTE3_EVENT);
|
||||
AscendC::DataCopyParams copyOutValueParams;
|
||||
copyOutValueParams.blockCount = 1;
|
||||
copyOutValueParams.blockLen = topkCount_ * sizeof(K_T); // bytes
|
||||
copyOutValueParams.srcStride = 0;
|
||||
copyOutValueParams.dstStride = 0;
|
||||
// 搬运到GM
|
||||
AscendC::DataCopyPad(
|
||||
valueOutGm[info.valueOutOffset + (curS1Idx + rowIdx) * topkCount_],
|
||||
valueOutLocal_, copyOutValueParams);
|
||||
}
|
||||
SetFlag<HardEvent::MTE3_V>(TOPK_MTE3_V_EVENT);
|
||||
}
|
||||
}
|
||||
} // namespace LIKernel
|
||||
#endif
|
||||
@@ -0,0 +1,181 @@
|
||||
/**
|
||||
* 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 lightning_indexer_topk.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef LIGHTNING_INDEXER_TOPK_H
|
||||
#define LIGHTNING_INDEXER_TOPK_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "vf_topk.h"
|
||||
#include "vf_topk_16_gather.h"
|
||||
|
||||
namespace topk {
|
||||
template<typename T>
|
||||
class LITopk {
|
||||
public:
|
||||
__aicore__ inline void operator()(LocalTensor<uint32_t>& outputIdxLocal,
|
||||
LocalTensor<T>& inputLocal,
|
||||
uint32_t s2SeqLen)
|
||||
{
|
||||
}
|
||||
};
|
||||
|
||||
template<>
|
||||
class LITopk<uint32_t> {
|
||||
public:
|
||||
static __aicore__ inline uint32_t GetSharedTmpBufferSize(uint32_t topK)
|
||||
{
|
||||
return 2 * topK * sizeof(uint32_t) + 5 * 256 * sizeof(uint32_t) + 64 * sizeof(uint32_t) +
|
||||
(topK + 64) * sizeof(uint32_t); // for output value tensor
|
||||
}
|
||||
|
||||
static __aicore__ inline uint32_t GetIndexBufferSize(uint32_t topK)
|
||||
{
|
||||
return (topK + 64) * sizeof(uint32_t);
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(uint32_t topK)
|
||||
{
|
||||
this->topK = topK;
|
||||
}
|
||||
|
||||
__aicore__ inline void InitBuffers(LocalTensor<uint32_t>& sharedTmpBuffer)
|
||||
{
|
||||
tmpIdxLocal = sharedTmpBuffer[0];
|
||||
tmpValueLocal = tmpIdxLocal[topK];
|
||||
histogramsLocal = tmpValueLocal[topK];
|
||||
idx0Local = histogramsLocal[256];
|
||||
idx1Local = idx0Local[256];
|
||||
idx2Local = idx1Local[256];
|
||||
idx3Local = idx2Local[256];
|
||||
nkValueLocal = idx3Local[256];
|
||||
outputValueLocal = nkValueLocal[64];
|
||||
}
|
||||
|
||||
__aicore__ inline void operator()(LocalTensor<uint32_t>& outputIdxLocal,
|
||||
LocalTensor<uint32_t>& inputLocal,
|
||||
uint32_t s2SeqLen)
|
||||
{
|
||||
topkb32::LiTopKVF(outputIdxLocal, // filter阶段使用输出value Buf topK * 4B
|
||||
outputValueLocal, // filter阶段使用输出 Idx Buf topK * 4B
|
||||
inputLocal, // 输入 s2SeqLen * 4B
|
||||
tmpIdxLocal, // filter阶段使用暂存index Buf topK * 4B
|
||||
tmpValueLocal, // filter阶段使用暂存value Buf topK * 4B
|
||||
histogramsLocal, // 直方图的临时Buf 256 * 4B
|
||||
idx0Local, // 输入数据第1个8位Buf 256 * 4B
|
||||
idx1Local, // 输入数据第2个8位Buf 256 * 4B
|
||||
idx2Local, // 输入数据第3个8位Buf 256 * 4B
|
||||
idx3Local, // 输入数据第4个8位Buf 256 * 4B
|
||||
nkValueLocal, // next_k 暂存Buf 64 * 4B
|
||||
topK, // topk数量
|
||||
s2SeqLen); // 输入元素总数
|
||||
}
|
||||
private:
|
||||
LocalTensor<uint32_t> tmpIdxLocal; // filter阶段使用暂存index Buf topK * 4B
|
||||
LocalTensor<uint32_t> tmpValueLocal; // filter阶段使用暂存value Buf topK * 4B
|
||||
LocalTensor<uint32_t> histogramsLocal; // 直方图的临时Buf 256 * 4B
|
||||
LocalTensor<uint32_t> idx0Local; // 输入数据第1个8位Buf 256 * 4B
|
||||
LocalTensor<uint32_t> idx1Local; // 输入数据第2个8位Buf 256 * 4B
|
||||
LocalTensor<uint32_t> idx2Local; // 输入数据第3个8位Buf 256 * 4B
|
||||
LocalTensor<uint32_t> idx3Local; // 输入数据第4个8位Buf 256 * 4B
|
||||
LocalTensor<uint32_t> nkValueLocal; // next_k 暂存Buf 64 * 4B
|
||||
LocalTensor<uint32_t> outputValueLocal; // 输出value tensor
|
||||
uint32_t topK;
|
||||
};
|
||||
|
||||
template<>
|
||||
class LITopk<uint16_t> {
|
||||
public:
|
||||
__aicore__ inline uint32_t GetSharedTmpBufferSize()
|
||||
{
|
||||
// 2 * LICommon::Align(topK, (uint32_t)256):两块hisIndexLocal；
|
||||
// 3 * 256：histogramsLocal idxHighLocal idxLowLocal；64：nkValueLocal
|
||||
uint64_t bufferSize1 = (2 * LICommon::Align(topK, (uint32_t)256) + 3 * 256 + 64) * sizeof(uint32_t);
|
||||
// LICommon::Align(topK, (uint32_t)256) + trunkLen：tmpIndexLocal
|
||||
uint64_t bufferSize2 = (LICommon::Align(topK, (uint32_t)256) + trunkLen) * sizeof(uint16_t);
|
||||
return bufferSize1 + bufferSize2;
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(uint32_t topK, uint32_t trunkLen)
|
||||
{
|
||||
this->topK = topK;
|
||||
this->trunkLen = trunkLen;
|
||||
}
|
||||
|
||||
__aicore__ inline void InitBuffers(LocalTensor<uint32_t>& sharedTmpBuffer)
|
||||
{
|
||||
LocalTensor<uint32_t> hisIndexLocal1 = sharedTmpBuffer[0];
|
||||
LocalTensor<uint32_t> hisIndexLocal2 = hisIndexLocal1[LICommon::Align(topK, (uint32_t)256)];
|
||||
hisIndexLocal[0] = hisIndexLocal1;
|
||||
hisIndexLocal[1] = hisIndexLocal2;
|
||||
histogramsLocal = hisIndexLocal2[LICommon::Align(topK, (uint32_t)256)];
|
||||
idxHighLocal = histogramsLocal[256];
|
||||
idxLowLocal = idxHighLocal[256];
|
||||
nkValueLocal = idxLowLocal[256];
|
||||
LocalTensor<uint32_t> tmpIndexLocalTmp = nkValueLocal[64];
|
||||
tmpIndexLocal = tmpIndexLocalTmp.template ReinterpretCast<uint16_t>();
|
||||
}
|
||||
|
||||
__aicore__ inline void operator()(LocalTensor<uint16_t>& mrgValueLocal, LocalTensor<uint32_t>& indicesOutLocal,
|
||||
LocalTensor<uint16_t>& hisValueLocal, uint32_t s2SeqLen, uint32_t loopIdx,
|
||||
uint32_t s2LoopNum, bool returnValueFlag)
|
||||
{
|
||||
if (s2LoopNum == 1) {
|
||||
if (returnValueFlag) {
|
||||
topkb16gather::LiTopKVF<true>(tmpIndexLocal, hisValueLocal,
|
||||
mrgValueLocal, histogramsLocal, idxHighLocal,
|
||||
idxLowLocal, nkValueLocal, topK, s2SeqLen);
|
||||
} else {
|
||||
topkb16gather::LiTopKVF<false>(tmpIndexLocal, hisValueLocal,
|
||||
mrgValueLocal, histogramsLocal, idxHighLocal,
|
||||
idxLowLocal, nkValueLocal, topK, s2SeqLen);
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(indicesOutLocal, tmpIndexLocal, RoundMode::CAST_NONE, topK);
|
||||
return;
|
||||
}
|
||||
|
||||
if (loopIdx == 0) {
|
||||
topkb16gather::LiTopKVF<true>(tmpIndexLocal, hisValueLocal,
|
||||
mrgValueLocal, histogramsLocal, idxHighLocal,
|
||||
idxLowLocal, nkValueLocal, topK, s2SeqLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Cast(hisIndexLocal[(loopIdx + 1) % 2], tmpIndexLocal, RoundMode::CAST_NONE, topK);
|
||||
} else {
|
||||
topkb16gather::LiTopKVF<true>(tmpIndexLocal, hisValueLocal,
|
||||
mrgValueLocal, histogramsLocal, idxHighLocal,
|
||||
idxLowLocal, nkValueLocal, topK, s2SeqLen);
|
||||
PipeBarrier<PIPE_V>();
|
||||
topkb16gather::LiTopKGatherVF(hisIndexLocal[(loopIdx + 1) % 2], hisValueLocal, mrgValueLocal, tmpIndexLocal,
|
||||
hisIndexLocal[loopIdx % 2], topK,
|
||||
loopIdx * trunkLen - LICommon::Align(topK, (uint32_t)256), s2SeqLen);
|
||||
if (loopIdx == s2LoopNum - 1) {
|
||||
PipeBarrier<PIPE_V>();
|
||||
AscendC::DataCopy(indicesOutLocal,
|
||||
hisIndexLocal[(loopIdx + 1) % 2],
|
||||
LICommon::Align(topK, (uint32_t)256));
|
||||
}
|
||||
}
|
||||
}
|
||||
private:
|
||||
LocalTensor<uint32_t> hisIndexLocal[2]; // 每trunkLen长度的s2选出的topK个索引
|
||||
LocalTensor<uint32_t> histogramsLocal; // 直方图的临时Buf 256 * 4B
|
||||
LocalTensor<uint32_t> idxHighLocal; // 输入数据高8位Buf 256 * 4B
|
||||
LocalTensor<uint32_t> idxLowLocal; // 输入数据低8位Buf 256 * 4B
|
||||
LocalTensor<uint32_t> nkValueLocal; // next_k 暂存Buf 64 * 4B
|
||||
LocalTensor<uint16_t> tmpIndexLocal; // 每trunkLen + topK的临时index
|
||||
uint32_t topK = 512;
|
||||
uint32_t trunkLen = 16384;
|
||||
};
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,830 @@
|
||||
/**
|
||||
* 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 lightning_indexer_vector1.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef LIGHTNING_INDEXER_VECTOR1_H
|
||||
#define LIGHTNING_INDEXER_VECTOR1_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
|
||||
namespace vector1 {
|
||||
|
||||
template <typename T>
|
||||
struct FloatSortTraits;
|
||||
|
||||
template <typename T>
|
||||
struct UIntSortTraits;
|
||||
|
||||
// fp32
|
||||
template <>
|
||||
struct FloatSortTraits<float> {
|
||||
using UInt = uint32_t;
|
||||
static constexpr UInt ZERO = 0x00000000;
|
||||
static constexpr UInt SIGN_MASK = 0x80000000;
|
||||
static constexpr UInt NAN_MASK = 0x7FC00000;
|
||||
static constexpr UInt ALL_ONE = 0xFFFFFFFF;
|
||||
};
|
||||
|
||||
// bf16
|
||||
template <>
|
||||
struct FloatSortTraits<bfloat16_t> {
|
||||
using UInt = uint16_t;
|
||||
static constexpr UInt ZERO = 0x0000;
|
||||
static constexpr UInt SIGN_MASK = 0x8000;
|
||||
static constexpr UInt NAN_MASK = 0x7FC0;
|
||||
static constexpr UInt ALL_ONE = 0xFFFF;
|
||||
};
|
||||
|
||||
|
||||
template <typename FloatT>
|
||||
struct FloatSortConstCtx {
|
||||
using Traits = FloatSortTraits<FloatT>;
|
||||
using UInt = typename Traits::UInt;
|
||||
AscendC::MicroAPI::RegTensor<UInt> zeros;
|
||||
AscendC::MicroAPI::RegTensor<UInt> all_one;
|
||||
AscendC::MicroAPI::RegTensor<UInt> signMask;
|
||||
AscendC::MicroAPI::RegTensor<UInt> nan;
|
||||
};
|
||||
|
||||
|
||||
template <typename FloatT>
|
||||
__simd_callee__ inline void InitFloatSortConstCtx(FloatSortConstCtx<FloatT>& ctx, AscendC::MicroAPI::MaskReg& maskAll)
|
||||
{
|
||||
using Traits = FloatSortTraits<FloatT>;
|
||||
AscendC::MicroAPI::Duplicate(ctx.zeros, Traits::ZERO, maskAll);
|
||||
AscendC::MicroAPI::Duplicate(ctx.all_one, Traits::ALL_ONE, maskAll);
|
||||
AscendC::MicroAPI::Duplicate(ctx.signMask, Traits::SIGN_MASK, maskAll);
|
||||
AscendC::MicroAPI::Duplicate(ctx.nan, Traits::NAN_MASK, maskAll);
|
||||
}
|
||||
|
||||
|
||||
template <typename FloatT>
|
||||
__simd_callee__ inline void FloatToSortableKey(
|
||||
AscendC::MicroAPI::RegTensor<typename FloatSortTraits<FloatT>::UInt>&
|
||||
outKey,
|
||||
AscendC::MicroAPI::RegTensor<FloatT>& inVal,
|
||||
FloatSortConstCtx<FloatT>& ctx,
|
||||
AscendC::MicroAPI::MaskReg& maskAll)
|
||||
{
|
||||
using Traits = FloatSortTraits<FloatT>;
|
||||
using UInt = typename Traits::UInt;
|
||||
|
||||
AscendC::MicroAPI::RegTensor<UInt> regTemp;
|
||||
AscendC::MicroAPI::RegTensor<UInt> regMask;
|
||||
AscendC::MicroAPI::MaskReg regSelectNan;
|
||||
AscendC::MicroAPI::MaskReg regSelectSign;
|
||||
|
||||
auto& inBits = (AscendC::MicroAPI::RegTensor<UInt>&)inVal;
|
||||
|
||||
// 1. NaN check
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::EQ>(regSelectNan, inBits, ctx.nan, maskAll);
|
||||
|
||||
// 2. NaN -> ALL_ONE
|
||||
AscendC::MicroAPI::Select(outKey, ctx.all_one, inBits, regSelectNan);
|
||||
|
||||
// 3. sign bit
|
||||
AscendC::MicroAPI::And(regTemp, outKey, ctx.signMask, maskAll);
|
||||
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::GT>(regSelectSign, regTemp, ctx.zeros, maskAll);
|
||||
|
||||
// 4. xor mask
|
||||
AscendC::MicroAPI::Select(regMask, ctx.all_one, ctx.signMask, regSelectSign);
|
||||
AscendC::MicroAPI::Xor(outKey, outKey, regMask, maskAll);
|
||||
}
|
||||
|
||||
// uint16-bf16
|
||||
template <>
|
||||
struct UIntSortTraits<bfloat16_t> {
|
||||
using UInt = uint16_t;
|
||||
static constexpr UInt ZERO = 0x0000;
|
||||
static constexpr UInt SIGN_MASK = 0x8000;
|
||||
static constexpr UInt NAN_MASK = 0xFFC0;
|
||||
static constexpr UInt ALL_ONE = 0xFFFF;
|
||||
};
|
||||
|
||||
template <typename FloatT>
|
||||
struct UIntSortConstCtx {
|
||||
using Traits = UIntSortTraits<FloatT>;
|
||||
using UInt = typename Traits::UInt;
|
||||
AscendC::MicroAPI::RegTensor<UInt> zeros;
|
||||
AscendC::MicroAPI::RegTensor<UInt> all_one;
|
||||
AscendC::MicroAPI::RegTensor<UInt> signMask;
|
||||
AscendC::MicroAPI::RegTensor<UInt> nan;
|
||||
};
|
||||
|
||||
template <typename FloatT>
|
||||
__simd_callee__ inline void InitUIntSortConstCtx(UIntSortConstCtx<FloatT>& ctx, AscendC::MicroAPI::MaskReg& maskAll)
|
||||
{
|
||||
using Traits = UIntSortTraits<FloatT>;
|
||||
AscendC::MicroAPI::Duplicate(ctx.zeros, Traits::ZERO, maskAll);
|
||||
AscendC::MicroAPI::Duplicate(ctx.all_one, Traits::ALL_ONE, maskAll);
|
||||
AscendC::MicroAPI::Duplicate(ctx.signMask, Traits::SIGN_MASK, maskAll);
|
||||
AscendC::MicroAPI::Duplicate(ctx.nan, Traits::NAN_MASK, maskAll);
|
||||
}
|
||||
|
||||
template <typename FloatT>
|
||||
__simd_callee__ inline void UIntToSortableKey(AscendC::MicroAPI::RegTensor<FloatT>& outKey,
|
||||
AscendC::MicroAPI::RegTensor<typename UIntSortConstCtx<FloatT>::UInt>&
|
||||
inVal,
|
||||
UIntSortConstCtx<FloatT>& ctx,
|
||||
AscendC::MicroAPI::MaskReg& maskAll)
|
||||
{
|
||||
using Traits = UIntSortTraits<FloatT>;
|
||||
using UInt = typename Traits::UInt;
|
||||
|
||||
AscendC::MicroAPI::RegTensor<UInt> regTemp;
|
||||
AscendC::MicroAPI::RegTensor<UInt> regMask;
|
||||
AscendC::MicroAPI::MaskReg regSelectZero;
|
||||
AscendC::MicroAPI::MaskReg regSelectSign;
|
||||
|
||||
auto& inBits = inVal;
|
||||
|
||||
// 1. 0 check
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::EQ>(regSelectZero, inBits, ctx.zeros, maskAll);
|
||||
|
||||
// 2. 0 -> -NAN
|
||||
AscendC::MicroAPI::Select((AscendC::MicroAPI::RegTensor<UInt>&)outKey, ctx.nan, inBits, regSelectZero);
|
||||
|
||||
// 3. sign bit
|
||||
AscendC::MicroAPI::And(regTemp, (AscendC::MicroAPI::RegTensor<UInt>&)outKey, ctx.signMask, maskAll);
|
||||
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::GT>(regSelectSign, regTemp, ctx.zeros, maskAll);
|
||||
|
||||
// 4. xor mask
|
||||
AscendC::MicroAPI::Select(regMask, ctx.signMask, ctx.all_one, regSelectSign);
|
||||
AscendC::MicroAPI::Xor((AscendC::MicroAPI::RegTensor<UInt>&)outKey,
|
||||
(AscendC::MicroAPI::RegTensor<UInt>&)outKey, regMask, maskAll);
|
||||
}
|
||||
|
||||
__aicore__ inline void UIntToFloatReturnValue(const LocalTensor<bfloat16_t> &out_,
|
||||
const LocalTensor<uint16_t> &in,
|
||||
const uint32_t topK)
|
||||
{
|
||||
auto outBuf = (__local_mem__ bfloat16_t*)out_.GetPhyAddr();
|
||||
auto inBuf = (__local_mem__ uint16_t*)in.GetPhyAddr();
|
||||
|
||||
const uint16_t repeatSize16 = 128;
|
||||
uint16_t topkLoopNum = (topK + repeatSize16 - 1) / repeatSize16;
|
||||
|
||||
__VEC_SCOPE__
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regIn;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regOut;
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
for (uint16_t i = 0; i < topkLoopNum; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<uint16_t>(regIn, inBuf + i * 128);
|
||||
|
||||
UIntSortConstCtx<bfloat16_t> uint16Ctx;
|
||||
InitUIntSortConstCtx(uint16Ctx, maskAllB16);
|
||||
|
||||
UIntToSortableKey<bfloat16_t>(regOut, regIn, uint16Ctx, maskAllB16);
|
||||
|
||||
AscendC::MicroAPI::StoreAlign<bfloat16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(
|
||||
outBuf + i * 128,
|
||||
regOut,
|
||||
maskAllB16);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void UIntToFloatReturnValue(const LocalTensor<half> &out_,
|
||||
const LocalTensor<uint16_t> &in,
|
||||
const uint32_t topK)
|
||||
{
|
||||
auto outBuf = (__local_mem__ half*)out_.GetPhyAddr();
|
||||
auto inBuf = (__local_mem__ uint16_t*)in.GetPhyAddr();
|
||||
|
||||
const uint16_t repeatSize16 = 128;
|
||||
uint16_t topkLoopNum = (topK + repeatSize16 - 1) / repeatSize16;
|
||||
|
||||
__VEC_SCOPE__
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regIn;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regOut;
|
||||
AscendC::MicroAPI::RegTensor<half> regOutHalf;
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllHalf =
|
||||
AscendC::MicroAPI::CreateMask<half, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
constexpr static MicroAPI::CastTrait castTraitBF16ToHalf =
|
||||
{MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_RINT};
|
||||
|
||||
for (uint16_t i = 0; i < topkLoopNum; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<uint16_t>(regIn, inBuf + i * repeatSize16);
|
||||
|
||||
UIntSortConstCtx<bfloat16_t> uint16Ctx;
|
||||
InitUIntSortConstCtx(uint16Ctx, maskAllB16);
|
||||
|
||||
UIntToSortableKey<bfloat16_t>(regOut, regIn, uint16Ctx, maskAllB16);
|
||||
|
||||
AscendC::MicroAPI::Cast<half, bfloat16_t, castTraitBF16ToHalf>(regOutHalf, regOut, maskAllB16);
|
||||
|
||||
AscendC::MicroAPI::StoreAlign<half, AscendC::MicroAPI::StoreDist::DIST_NORM>(
|
||||
outBuf + i * repeatSize16,
|
||||
regOutHalf,
|
||||
maskAllHalf);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename FloatT>
|
||||
__simd_callee__ inline void FloatX2ToSortableKey(AscendC::MicroAPI::RegTensor<typename FloatSortTraits<FloatT>::UInt>&
|
||||
outKey0,
|
||||
AscendC::MicroAPI::RegTensor<typename FloatSortTraits<FloatT>::UInt>&
|
||||
outKey1,
|
||||
AscendC::MicroAPI::RegTensor<FloatT>& inVal0,
|
||||
AscendC::MicroAPI::RegTensor<FloatT>& inVal1,
|
||||
FloatSortConstCtx<FloatT>& ctx,
|
||||
AscendC::MicroAPI::MaskReg& maskAll)
|
||||
{
|
||||
using Traits = FloatSortTraits<FloatT>;
|
||||
using UInt = typename Traits::UInt;
|
||||
|
||||
AscendC::MicroAPI::RegTensor<UInt> regTemp[2];
|
||||
AscendC::MicroAPI::RegTensor<UInt> regMask[2];
|
||||
AscendC::MicroAPI::MaskReg regSelectNan[2];
|
||||
AscendC::MicroAPI::MaskReg regSelectSign[2];
|
||||
|
||||
auto& inBits0 = (AscendC::MicroAPI::RegTensor<UInt>&)inVal0;
|
||||
auto& inBits1 = (AscendC::MicroAPI::RegTensor<UInt>&)inVal1;
|
||||
|
||||
// 1. NaN check
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::EQ>(regSelectNan[0], inBits0, ctx.nan, maskAll);
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::EQ>(regSelectNan[1], inBits1, ctx.nan, maskAll);
|
||||
|
||||
// 2. NaN -> ALL_ONE
|
||||
AscendC::MicroAPI::Select(outKey0, ctx.all_one, inBits0, regSelectNan[0]);
|
||||
AscendC::MicroAPI::Select(outKey1, ctx.all_one, inBits1, regSelectNan[1]);
|
||||
|
||||
// 3. sign bit
|
||||
AscendC::MicroAPI::And(regTemp[0], outKey0, ctx.signMask, maskAll);
|
||||
AscendC::MicroAPI::And(regTemp[1], outKey1, ctx.signMask, maskAll);
|
||||
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::GT>(regSelectSign[0], regTemp[0], ctx.zeros, maskAll);
|
||||
AscendC::MicroAPI::Compare<UInt, CMPMODE::GT>(regSelectSign[1], regTemp[1], ctx.zeros, maskAll);
|
||||
|
||||
// 4. xor mask
|
||||
AscendC::MicroAPI::Select(regMask[0], ctx.all_one, ctx.signMask, regSelectSign[0]);
|
||||
AscendC::MicroAPI::Select(regMask[1], ctx.all_one, ctx.signMask, regSelectSign[1]);
|
||||
AscendC::MicroAPI::Xor(outKey0, outKey0, regMask[0], maskAll);
|
||||
AscendC::MicroAPI::Xor(outKey1, outKey1, regMask[1], maskAll);
|
||||
}
|
||||
|
||||
|
||||
template <typename T, size_t N>
|
||||
__simd_callee__ inline void DuplicateZero(AscendC::MicroAPI::RegTensor<T> (®Array)[N],
|
||||
AscendC::MicroAPI::MaskReg& mask)
|
||||
{
|
||||
static_assert(N <= 4, "N must be <= 4");
|
||||
// 不能用循环, 会导致fatal error: error in backend: Unsupported Inst must be hoisted.
|
||||
if constexpr (N >= 1) {
|
||||
AscendC::MicroAPI::Duplicate(regArray[0], static_cast<T>(0), mask);
|
||||
}
|
||||
if constexpr (N >= 2) {
|
||||
AscendC::MicroAPI::Duplicate(regArray[1], static_cast<T>(0), mask);
|
||||
}
|
||||
if constexpr (N >= 3) {
|
||||
AscendC::MicroAPI::Duplicate(regArray[2], static_cast<T>(0), mask);
|
||||
}
|
||||
if constexpr (N >= 4) {
|
||||
AscendC::MicroAPI::Duplicate(regArray[3], static_cast<T>(0), mask);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T, size_t N, bool ApplyRelu = true>
|
||||
__simd_callee__ inline void WeightedAccum(AscendC::MicroAPI::RegTensor<T> (&accum)[N],
|
||||
AscendC::MicroAPI::RegTensor<T> (&input)[N],
|
||||
AscendC::MicroAPI::RegTensor<T>& weight,
|
||||
AscendC::MicroAPI::MaskReg& mask)
|
||||
{
|
||||
static_assert(N <= 2, "N must be <= 2");
|
||||
// ---- Relu block ----
|
||||
if constexpr (ApplyRelu) {
|
||||
if constexpr (N >= 1) {
|
||||
AscendC::MicroAPI::Relu(input[0], input[0], mask);
|
||||
}
|
||||
if constexpr (N >= 2) {
|
||||
AscendC::MicroAPI::Relu(input[1], input[1], mask);
|
||||
}
|
||||
}
|
||||
// ---- MulAdd block ----
|
||||
if constexpr (N >= 1) {
|
||||
AscendC::MicroAPI::MulAddDst(accum[0], input[0], weight, mask);
|
||||
}
|
||||
if constexpr (N >= 2) {
|
||||
AscendC::MicroAPI::MulAddDst(accum[1], input[1], weight, mask);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
__simd_callee__ inline void BroadcastLane(AscendC::MicroAPI::RegTensor<float>& dst,
|
||||
AscendC::MicroAPI::RegTensor<float>& src,
|
||||
uint16_t laneIdx)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<uint32_t> brcGatherIndex;
|
||||
AscendC::MicroAPI::Duplicate(brcGatherIndex, laneIdx);
|
||||
AscendC::MicroAPI::Gather(dst, src, brcGatherIndex);
|
||||
}
|
||||
|
||||
__simd_callee__ inline void BroadcastLane(AscendC::MicroAPI::RegTensor<float>& dst,
|
||||
__local_mem__ float* src,
|
||||
uint16_t laneIdx)
|
||||
{
|
||||
AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(dst, src + laneIdx);
|
||||
}
|
||||
|
||||
// float in uint16 out
|
||||
__simd_vf__ inline void MulWeightAndReduceSum(__ubuf__ uint16_t* out_,
|
||||
__ubuf__ float* qk_,
|
||||
const uint32_t qkVLStride,
|
||||
__ubuf__ float* weight_,
|
||||
const int gSize)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> regwBrc;
|
||||
AscendC::MicroAPI::RegTensor<float> regQK[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regW;
|
||||
|
||||
AscendC::MicroAPI::RegTensor<float> regSum0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regSum1[2];
|
||||
AscendC::MicroAPI::MaskReg maskAllB32 =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
FloatSortConstCtx<bfloat16_t> bf16Ctx;
|
||||
InitFloatSortConstCtx(bf16Ctx, maskAllB16);
|
||||
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN =
|
||||
{MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::MERGING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD =
|
||||
{MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_ROUND};
|
||||
|
||||
AscendC::MicroAPI::LoadAlign<float>(regW, weight_);
|
||||
DuplicateZero(regSum0, maskAllB32);
|
||||
DuplicateZero(regSum1, maskAllB32);
|
||||
|
||||
// unroll2
|
||||
for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) {
|
||||
MicroAPI::LoadAlign<float>(regQK[0], qk_ + 128 * i); // RowStride是128, 行都落在一个bank上
|
||||
MicroAPI::LoadAlign<float>(regQK[1], qk_ + 128 * i + qkVLStride);
|
||||
BroadcastLane(regwBrc, regW, i);
|
||||
WeightedAccum(regSum0, regQK, regwBrc, maskAllB32);
|
||||
|
||||
MicroAPI::LoadAlign<float>(regQK[0], qk_ + 128 * i + 128);
|
||||
MicroAPI::LoadAlign<float>(regQK[1], qk_ + 128 * i + 128 + qkVLStride);
|
||||
BroadcastLane(regwBrc, regW, i + 1);
|
||||
WeightedAccum(regSum1, regQK, regwBrc, maskAllB32);
|
||||
}
|
||||
|
||||
AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32);
|
||||
AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32);
|
||||
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regSumBF16;
|
||||
// interleave cast ==> regSum[1] high regSum[0] low
|
||||
AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16, regSum0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16, regSum0[0], maskAllB32);
|
||||
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regOut;
|
||||
FloatToSortableKey<bfloat16_t>(regOut, regSumBF16, bf16Ctx, maskAllB16);
|
||||
// normal store
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out_, regOut, maskAllB16);
|
||||
}
|
||||
|
||||
// float in uint16 out
|
||||
__simd_vf__ inline void MulWeightAndReduceSum(__ubuf__ uint16_t* out_,
|
||||
__ubuf__ float* qk_,
|
||||
const uint32_t qkVLStride,
|
||||
__ubuf__ bfloat16_t* weight_,
|
||||
const int gSize)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> regwBrc;
|
||||
AscendC::MicroAPI::RegTensor<float> regQK[2];
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regWBF16;
|
||||
AscendC::MicroAPI::RegTensor<float> regW;
|
||||
|
||||
AscendC::MicroAPI::RegTensor<float> regSum0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regSum1[2];
|
||||
AscendC::MicroAPI::MaskReg maskAllB32 =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
FloatSortConstCtx<bfloat16_t> bf16Ctx;
|
||||
InitFloatSortConstCtx(bf16Ctx, maskAllB16);
|
||||
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::MERGING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitBF16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::UNKNOWN};
|
||||
|
||||
AscendC::MicroAPI::LoadAlign<bfloat16_t, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regWBF16, weight_);
|
||||
AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitBF16ToFP32>(regW, regWBF16, maskAllB16);
|
||||
|
||||
DuplicateZero(regSum0, maskAllB32);
|
||||
DuplicateZero(regSum1, maskAllB32);
|
||||
|
||||
// unroll2
|
||||
for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) {
|
||||
MicroAPI::LoadAlign<float>(regQK[0], qk_ + 128 * i); // RowStride是128, 行都落在一个bank上
|
||||
MicroAPI::LoadAlign<float>(regQK[1], qk_ + 128 * i + qkVLStride);
|
||||
BroadcastLane(regwBrc, regW, i);
|
||||
WeightedAccum(regSum0, regQK, regwBrc, maskAllB32);
|
||||
|
||||
MicroAPI::LoadAlign<float>(regQK[0], qk_ + 128 * i + 128);
|
||||
MicroAPI::LoadAlign<float>(regQK[1], qk_ + 128 * i + 128 + qkVLStride);
|
||||
BroadcastLane(regwBrc, regW, i + 1);
|
||||
WeightedAccum(regSum1, regQK, regwBrc, maskAllB32);
|
||||
}
|
||||
|
||||
AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32);
|
||||
AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32);
|
||||
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regSumBF16;
|
||||
// interleave cast ==> regSum[1] high regSum[0] low
|
||||
AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16, regSum0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16, regSum0[0], maskAllB32);
|
||||
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regOut;
|
||||
FloatToSortableKey<bfloat16_t>(regOut, regSumBF16, bf16Ctx, maskAllB16);
|
||||
// normal store
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out_, regOut, maskAllB16);
|
||||
}
|
||||
|
||||
// float in uint16 out
|
||||
__simd_vf__ inline void MulWeightAndReduceSum(__ubuf__ uint16_t* out_,
|
||||
__ubuf__ float* qk_,
|
||||
const uint32_t qkVLStride,
|
||||
__ubuf__ half* weight_,
|
||||
const int gSize)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> regwBrc;
|
||||
AscendC::MicroAPI::RegTensor<float> regQK[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regW;
|
||||
AscendC::MicroAPI::RegTensor<half> regWFP16;
|
||||
AscendC::MicroAPI::RegTensor<float> regSum0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regSum1[2];
|
||||
AscendC::MicroAPI::MaskReg maskAllB32 =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
FloatSortConstCtx<bfloat16_t> bf16Ctx;
|
||||
InitFloatSortConstCtx(bf16Ctx, maskAllB16);
|
||||
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::MERGING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitFP16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::UNKNOWN};
|
||||
|
||||
AscendC::MicroAPI::LoadAlign<half, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regWFP16, weight_);
|
||||
AscendC::MicroAPI::Cast<float, half, castTraitFP16ToFP32>(regW, regWFP16, maskAllB16);
|
||||
|
||||
DuplicateZero(regSum0, maskAllB32);
|
||||
DuplicateZero(regSum1, maskAllB32);
|
||||
|
||||
// unroll2
|
||||
for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i += 2) {
|
||||
MicroAPI::LoadAlign<float>(regQK[0], qk_ + 128 * i); // RowStride是128, 行都落在一个bank上
|
||||
MicroAPI::LoadAlign<float>(regQK[1], qk_ + 128 * i + qkVLStride);
|
||||
BroadcastLane(regwBrc, regW, i);
|
||||
WeightedAccum(regSum0, regQK, regwBrc, maskAllB32);
|
||||
|
||||
MicroAPI::LoadAlign<float>(regQK[0], qk_ + 128 * i + 128);
|
||||
MicroAPI::LoadAlign<float>(regQK[1], qk_ + 128 * i + 128 + qkVLStride);
|
||||
BroadcastLane(regwBrc, regW, i + 1);
|
||||
WeightedAccum(regSum1, regQK, regwBrc, maskAllB32);
|
||||
}
|
||||
|
||||
AscendC::MicroAPI::Add(regSum0[0], regSum0[0], regSum1[0], maskAllB32);
|
||||
AscendC::MicroAPI::Add(regSum0[1], regSum0[1], regSum1[1], maskAllB32);
|
||||
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regSumBF16;
|
||||
// interleave cast ==> regSum[1] high regSum[0] low
|
||||
AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16, regSum0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16, regSum0[0], maskAllB32);
|
||||
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regOut;
|
||||
FloatToSortableKey<bfloat16_t>(regOut, regSumBF16, bf16Ctx, maskAllB16);
|
||||
// normal store
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out_, regOut, maskAllB16);
|
||||
}
|
||||
|
||||
// 计算S1=2
|
||||
// float in uint16 out
|
||||
__simd_vf__ inline void MulWeightAndReduceSum2(__ubuf__ uint16_t* out0_,
|
||||
__ubuf__ uint16_t* out1_,
|
||||
uint32_t outStride,
|
||||
__ubuf__ float* qk0_,
|
||||
__ubuf__ float* qk1_,
|
||||
uint32_t qkVLStride,
|
||||
uint32_t qkStride,
|
||||
__ubuf__ float* weight0_,
|
||||
__ubuf__ float* weight1_,
|
||||
uint32_t weightStride,
|
||||
__ubuf__ float* weightFloat_,
|
||||
const int gSize)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> regwBrc[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regQK0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regQK1[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regW[2];
|
||||
|
||||
AscendC::MicroAPI::RegTensor<float> regSum0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regSum1[2];
|
||||
AscendC::MicroAPI::MaskReg maskAllB32 =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
FloatSortConstCtx<bfloat16_t> bf16Ctx;
|
||||
InitFloatSortConstCtx(bf16Ctx, maskAllB16);
|
||||
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::MERGING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_ROUND};
|
||||
|
||||
AscendC::MicroAPI::LoadAlign<float>(regW[0], weight0_);
|
||||
AscendC::MicroAPI::LoadAlign<float>(regW[1], weight1_);
|
||||
// regW[0]与weight1混合使用
|
||||
AscendC::MicroAPI::StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_NORM>(weight1_, regW[1], maskAllB32);
|
||||
AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
DuplicateZero(regSum0, maskAllB32);
|
||||
DuplicateZero(regSum1, maskAllB32);
|
||||
|
||||
for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) {
|
||||
MicroAPI::LoadAlign<float>(regQK0[0], qk0_ + 128 * i);
|
||||
MicroAPI::LoadAlign<float>(regQK0[1], qk0_ + 128 * i + qkVLStride);
|
||||
MicroAPI::LoadAlign<float>(regQK1[0], qk1_ + 128 * i);
|
||||
MicroAPI::LoadAlign<float>(regQK1[1], qk1_ + 128 * i + qkVLStride);
|
||||
// 混合使用对整体性能更好
|
||||
BroadcastLane(regwBrc[0], regW[0], i);
|
||||
// Weight无bank冲突，用LoadAlign来提取weight标量
|
||||
BroadcastLane(regwBrc[1], weight1_, i);
|
||||
AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32);
|
||||
}
|
||||
|
||||
// Convert to bfloat16 and store output channel
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regSumBF16[2];
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regOut[2];
|
||||
AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]);
|
||||
AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16[0], regSum0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16[1], regSum1[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16[0], regSum0[0], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16[1], regSum1[0], maskAllB32);
|
||||
|
||||
FloatX2ToSortableKey<bfloat16_t>(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16);
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out0_, regOut[0], maskAllB16);
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out1_, regOut[1], maskAllB16);
|
||||
}
|
||||
|
||||
// 计算S1=2
|
||||
// float in uint16 out
|
||||
__simd_vf__ inline void MulWeightAndReduceSum2(__ubuf__ uint16_t* out0_,
|
||||
__ubuf__ uint16_t* out1_,
|
||||
uint32_t outStride,
|
||||
__ubuf__ float* qk0_,
|
||||
__ubuf__ float* qk1_,
|
||||
uint32_t qkVLStride,
|
||||
uint32_t qkStride,
|
||||
__ubuf__ bfloat16_t* weight0_,
|
||||
__ubuf__ bfloat16_t* weight1_,
|
||||
uint32_t weightStride,
|
||||
__ubuf__ float* weightFloat_,
|
||||
const int gSize)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> regwBrc[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regQK0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regQK1[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regW[2];
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regWBF16[2];
|
||||
|
||||
AscendC::MicroAPI::RegTensor<float> regSum0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regSum1[2];
|
||||
AscendC::MicroAPI::MaskReg maskAllB32 =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
FloatSortConstCtx<bfloat16_t> bf16Ctx;
|
||||
InitFloatSortConstCtx(bf16Ctx, maskAllB16);
|
||||
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::MERGING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitBF16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::UNKNOWN};
|
||||
|
||||
AscendC::MicroAPI::LoadAlign<bfloat16_t, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regWBF16[0], weight0_);
|
||||
AscendC::MicroAPI::LoadAlign<bfloat16_t, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regWBF16[1], weight1_);
|
||||
AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitBF16ToFP32>(regW[0], regWBF16[0], maskAllB16);
|
||||
AscendC::MicroAPI::Cast<float, bfloat16_t, castTraitBF16ToFP32>(regW[1], regWBF16[1], maskAllB16);
|
||||
|
||||
// regW[0]与weight1混合使用
|
||||
AscendC::MicroAPI::StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_NORM>(weightFloat_, regW[1], maskAllB32);
|
||||
AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
DuplicateZero(regSum0, maskAllB32);
|
||||
DuplicateZero(regSum1, maskAllB32);
|
||||
|
||||
for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) {
|
||||
MicroAPI::LoadAlign<float>(regQK0[0], qk0_ + 128 * i);
|
||||
MicroAPI::LoadAlign<float>(regQK0[1], qk0_ + 128 * i + qkVLStride);
|
||||
MicroAPI::LoadAlign<float>(regQK1[0], qk1_ + 128 * i);
|
||||
MicroAPI::LoadAlign<float>(regQK1[1], qk1_ + 128 * i + qkVLStride);
|
||||
// 混合使用对整体性能更好
|
||||
BroadcastLane(regwBrc[0], regW[0], i);
|
||||
// Weight无bank冲突，用LoadAlign来提取weight标量
|
||||
BroadcastLane(regwBrc[1], weightFloat_, i);
|
||||
AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32);
|
||||
}
|
||||
|
||||
// Convert to bfloat16 and store output channel
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regSumBF16[2];
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regOut[2];
|
||||
AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]);
|
||||
AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16[0], regSum0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16[1], regSum1[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16[0], regSum0[0], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16[1], regSum1[0], maskAllB32);
|
||||
|
||||
FloatX2ToSortableKey<bfloat16_t>(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16);
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out0_, regOut[0], maskAllB16);
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out1_, regOut[1], maskAllB16);
|
||||
}
|
||||
|
||||
// 计算S1=2
|
||||
// float in uint16 out
|
||||
__simd_vf__ inline void MulWeightAndReduceSum2(__ubuf__ uint16_t* out0_,
|
||||
__ubuf__ uint16_t* out1_,
|
||||
uint32_t outStride,
|
||||
__ubuf__ float* qk0_,
|
||||
__ubuf__ float* qk1_,
|
||||
uint32_t qkVLStride,
|
||||
uint32_t qkStride,
|
||||
__ubuf__ half* weight0_,
|
||||
__ubuf__ half* weight1_,
|
||||
uint32_t weightStride,
|
||||
__ubuf__ float* weightFloat_,
|
||||
const int gSize)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> regwBrc[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regQK0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regQK1[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regW[2];
|
||||
AscendC::MicroAPI::RegTensor<half> regWFP16[2];
|
||||
|
||||
AscendC::MicroAPI::RegTensor<float> regSum0[2];
|
||||
AscendC::MicroAPI::RegTensor<float> regSum1[2];
|
||||
AscendC::MicroAPI::MaskReg maskAllB32 =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg maskAllB16 =
|
||||
AscendC::MicroAPI::CreateMask<bfloat16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
FloatSortConstCtx<bfloat16_t> bf16Ctx;
|
||||
InitFloatSortConstCtx(bf16Ctx, maskAllB16);
|
||||
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_EVEN = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::MERGING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitF32ToF16_ODD = {MicroAPI::RegLayout::ONE, MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_ROUND};
|
||||
constexpr static MicroAPI::CastTrait castTraitFP16ToFP32 = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::UNKNOWN};
|
||||
|
||||
AscendC::MicroAPI::LoadAlign<half, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regWFP16[0], weight0_);
|
||||
AscendC::MicroAPI::LoadAlign<half, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regWFP16[1], weight1_);
|
||||
AscendC::MicroAPI::Cast<float, half, castTraitFP16ToFP32>(regW[0], regWFP16[0], maskAllB16);
|
||||
AscendC::MicroAPI::Cast<float, half, castTraitFP16ToFP32>(regW[1], regWFP16[1], maskAllB16);
|
||||
|
||||
// regW[0]与weight1混合使用
|
||||
AscendC::MicroAPI::StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_NORM>(weightFloat_, regW[1], maskAllB32);
|
||||
AscendC::MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
DuplicateZero(regSum0, maskAllB32);
|
||||
DuplicateZero(regSum1, maskAllB32);
|
||||
|
||||
for (uint16_t i = (uint16_t)(0); i < (uint16_t)(gSize); i++) {
|
||||
MicroAPI::LoadAlign<float>(regQK0[0], qk0_ + 128 * i);
|
||||
MicroAPI::LoadAlign<float>(regQK0[1], qk0_ + 128 * i + qkVLStride);
|
||||
MicroAPI::LoadAlign<float>(regQK1[0], qk1_ + 128 * i);
|
||||
MicroAPI::LoadAlign<float>(regQK1[1], qk1_ + 128 * i + qkVLStride);
|
||||
// 混合使用对整体性能更好
|
||||
BroadcastLane(regwBrc[0], regW[0], i);
|
||||
// Weight无bank冲突，用LoadAlign来提取weight标量
|
||||
BroadcastLane(regwBrc[1], weightFloat_, i);
|
||||
AscendC::MicroAPI::Relu(regQK0[0], regQK0[0], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK0[1], regQK0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK1[0], regQK1[0], maskAllB32);
|
||||
AscendC::MicroAPI::Relu(regQK1[1], regQK1[1], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum0[0], regQK0[0], regwBrc[0], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum0[1], regQK0[1], regwBrc[0], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum1[0], regQK1[0], regwBrc[1], maskAllB32);
|
||||
AscendC::MicroAPI::MulAddDst(regSum1[1], regQK1[1], regwBrc[1], maskAllB32);
|
||||
}
|
||||
|
||||
// Convert to bfloat16 and store output channel
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> regSumBF16[2];
|
||||
AscendC::MicroAPI::RegTensor<uint16_t> regOut[2];
|
||||
AscendC::MicroAPI::DeInterleave(regSum0[0], regSum0[1], regSum0[0], regSum0[1]);
|
||||
AscendC::MicroAPI::DeInterleave(regSum1[0], regSum1[1], regSum1[0], regSum1[1]);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16[0], regSum0[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_ODD>(regSumBF16[1], regSum1[1], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16[0], regSum0[0], maskAllB32);
|
||||
AscendC::MicroAPI::Cast<bfloat16_t, float, castTraitF32ToF16_EVEN>(regSumBF16[1], regSum1[0], maskAllB32);
|
||||
|
||||
FloatX2ToSortableKey<bfloat16_t>(regOut[0], regOut[1], regSumBF16[0], regSumBF16[1], bf16Ctx, maskAllB16);
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out0_, regOut[0], maskAllB16);
|
||||
AscendC::MicroAPI::StoreAlign<uint16_t, AscendC::MicroAPI::StoreDist::DIST_NORM>(out1_, regOut[1], maskAllB16);
|
||||
}
|
||||
|
||||
template<typename QK_T, typename W_T, typename SCORE_T>
|
||||
__aicore__ inline void BatchMulWeightAndReduceSum(const LocalTensor<SCORE_T> &out_, // out [S2Base] [128 ]
|
||||
uint32_t outStride,
|
||||
const LocalTensor<QK_T> &qk_, // q*k^t [G, S2Base] [64 128]
|
||||
uint32_t qkVLStride,
|
||||
uint32_t qkStride,
|
||||
const LocalTensor<W_T> &weight_, // w [G] [64 ]
|
||||
uint32_t weightStride,
|
||||
const LocalTensor<float> &weightFloat_,
|
||||
const int gSize, // G 64
|
||||
const int batch)
|
||||
{
|
||||
// 暂只支持这两种情况, 后续改成循环
|
||||
if (batch != 2 && batch != 1) {
|
||||
return;
|
||||
}
|
||||
auto weight = (__ubuf__ W_T *)weight_.GetPhyAddr();
|
||||
auto weightFloat = (__ubuf__ float *)weightFloat_.GetPhyAddr();
|
||||
auto qk = (__ubuf__ float *)qk_.GetPhyAddr();
|
||||
auto out = (__ubuf__ uint16_t *)out_.GetPhyAddr();
|
||||
if (batch == 2) {
|
||||
auto weight1 = weight + weightStride;
|
||||
auto qk1 = qk + qkStride;
|
||||
auto out1 = out + outStride;
|
||||
MulWeightAndReduceSum2(out, out1, outStride,
|
||||
qk, qk1, qkVLStride, qkStride,
|
||||
weight, weight1, weightStride, weightFloat,
|
||||
gSize);
|
||||
} else {
|
||||
MulWeightAndReduceSum(out, qk, qkVLStride, weight, gSize);
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#endif
|
||||
736
csrc/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk.h
Normal file
736
csrc/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk.h
Normal file
@@ -0,0 +1,736 @@
|
||||
/**
|
||||
* 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 vf_top_k.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef VF_TOP_K_H
|
||||
#define VF_TOP_K_H
|
||||
|
||||
namespace topkb32 {
|
||||
template<typename T>
|
||||
__simd_vf__ void HistogramsFirstVFImpl(__ubuf__ uint32_t* histogramsBuf,
|
||||
__ubuf__ uint32_t* inputBuf,
|
||||
uint16_t vfLoop, bool init)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
// 计算直方图cout0 0-127 cout1 128-255
|
||||
MicroAPI::RegTensor<uint16_t> cout0;
|
||||
MicroAPI::RegTensor<uint16_t> cout1;
|
||||
MicroAPI::Duplicate(cout0, 0);
|
||||
MicroAPI::Duplicate(cout1, 0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Odd;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Odd;
|
||||
|
||||
// 32bit 高16bit
|
||||
MicroAPI::RegTensor<uint32_t> vreg0U16;
|
||||
// 32bit 低16bit
|
||||
MicroAPI::RegTensor<uint32_t> vreg1U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg2U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg3U16;
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> vreg0;
|
||||
MicroAPI::RegTensor<uint8_t> vreg1;
|
||||
MicroAPI::RegTensor<uint8_t> vreg2;
|
||||
MicroAPI::RegTensor<uint8_t> vreg3;
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
for (uint16_t i = 0; i < vfLoop; ++i) {
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg1U16, vreg0U16, inputBuf + i * 256);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(
|
||||
vreg3U16, vreg2U16, inputBuf + (i * 256) + 128);
|
||||
|
||||
MicroAPI::DeInterleave(vreg1, vreg0,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg0U16,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg2U16);
|
||||
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout0, vreg0, pregB8);
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout1, vreg0, pregB8);
|
||||
}
|
||||
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout0U32Even, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout0U32Odd, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout1U32Even, cout1, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout1U32Odd, cout1, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
|
||||
histogramsBuf, cout0U32Even, cout0U32Odd, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(
|
||||
histogramsBuf + 128, cout1U32Even, cout1U32Odd, pregB32);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindFirstTargetBinVFImpl(__ubuf__ uint32_t* idx0Buf,
|
||||
__ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t*
|
||||
histogramsBuf, uint32_t bottomK)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx0;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> btmK;
|
||||
MicroAPI::Duplicate(btmK, bottomK);
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(4); ++i) {
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::RegTensor<uint32_t> cout;
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdx0;
|
||||
|
||||
MicroAPI::MaskReg pregGE = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::Arange(idxC, i * 64);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(cout, histogramsBuf + i * 64);
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GE>(pregGE, cout, btmK, pregB32);
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(
|
||||
sqzIdx0, (MicroAPI::RegTensor<uint32_t>&)idxC, pregGE);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(idx0Buf, sqzIdx0, alignIdx0);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(idx0Buf, alignIdx0);
|
||||
|
||||
MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx0;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx0, idx0Buf);
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> idxAll1;
|
||||
MicroAPI::RegTensor<uint32_t> idxPrev0;
|
||||
MicroAPI::RegTensor<uint32_t> prevBinValue;
|
||||
MicroAPI::Duplicate(idxAll1, 1);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> zeroAll;
|
||||
MicroAPI::Duplicate(zeroAll, 0);
|
||||
|
||||
MicroAPI::MaskReg preg0 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::EQ>(preg0, idx0, zeroAll, pregB32);
|
||||
MicroAPI::Sub(idxPrev0, idx0, (MicroAPI::RegTensor<uint32_t>&)idxAll1, pregB32);
|
||||
MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32);
|
||||
|
||||
MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32);
|
||||
MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> nextK;
|
||||
MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM>(nkValueBuf, nextK, pregB32);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__simd_vf__ void HistogramsSecondVFImpl(__ubuf__ uint32_t* histogramsBuf,
|
||||
__ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf,
|
||||
uint16_t vfLoop, bool init)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
// 计算直方图0-127 128-255
|
||||
MicroAPI::RegTensor<uint16_t> cout0;
|
||||
MicroAPI::RegTensor<uint16_t> cout1;
|
||||
MicroAPI::Duplicate(cout0, 0);
|
||||
MicroAPI::Duplicate(cout1, 0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Odd;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Odd;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx0;
|
||||
// 0x000000fc -> 0xfcfcfcfc
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx0, idx0Buf);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vreg0U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg1U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg2U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg3U16;
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> vreg0;
|
||||
MicroAPI::RegTensor<uint8_t> vreg1;
|
||||
MicroAPI::RegTensor<uint8_t> vreg2;
|
||||
MicroAPI::RegTensor<uint8_t> vreg3;
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
for (uint16_t i = 0; i < vfLoop; ++i) {
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg1U16,
|
||||
vreg0U16, inputBuf + i * 256);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg3U16,
|
||||
vreg2U16, inputBuf + (i * 256) + 128);
|
||||
|
||||
MicroAPI::DeInterleave(vreg1, vreg0,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg0U16,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg2U16);
|
||||
|
||||
MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ, vreg0, (MicroAPI::RegTensor<uint8_t>&)idx0, pregB8);
|
||||
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout0, vreg1, pregEQ);
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout1, vreg1, pregEQ);
|
||||
}
|
||||
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout0U32Even, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout0U32Odd, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout1U32Even, cout1, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout1U32Odd, cout1, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf,
|
||||
cout0U32Even, cout0U32Odd, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf + 128,
|
||||
cout1U32Even, cout1U32Odd, pregB32);
|
||||
}
|
||||
|
||||
// kValue新的bottomK
|
||||
__simd_vf__ void FindSecondTargetBinVFImpl(__ubuf__ uint32_t* idx1Buf,
|
||||
__ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue,
|
||||
__ubuf__ uint32_t* histogramsBuf)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx1;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> btmK1;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(btmK1, kValue);
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(4); ++i) {
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::RegTensor<uint32_t> cout;
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdx1;
|
||||
|
||||
MicroAPI::MaskReg pregGE = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::Arange(idxC, i * 64);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(cout, histogramsBuf + i * 64);
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GE>(pregGE, cout, btmK1, pregB32);
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdx1,
|
||||
(MicroAPI::RegTensor<uint32_t>&)idxC, pregGE);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(idx1Buf, sqzIdx1, alignIdx1);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(idx1Buf, alignIdx1);
|
||||
|
||||
MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx1;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx1, idx1Buf);
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> idxAll1;
|
||||
MicroAPI::RegTensor<uint32_t> idxPrev1;
|
||||
MicroAPI::RegTensor<uint32_t> prevBinValue;
|
||||
MicroAPI::Duplicate(idxAll1, 1);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> zeroAll;
|
||||
MicroAPI::Duplicate(zeroAll, 0);
|
||||
|
||||
MicroAPI::MaskReg preg1 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::EQ>(preg1, idx1, zeroAll, pregB32);
|
||||
MicroAPI::Sub(idxPrev1, idx1, (MicroAPI::RegTensor<uint32_t>&)idxAll1, pregB32);
|
||||
MicroAPI::ShiftRights(idxPrev1, idxPrev1, (int16_t)24, pregB32);
|
||||
|
||||
MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev1, pregB32);
|
||||
MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg1);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> nextK;
|
||||
MicroAPI::Sub(nextK, btmK1, prevBinValue, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM>(nkValueBuf, nextK, pregB32);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__simd_vf__ void HistogramsThirdVFImpl(__ubuf__ uint32_t* histogramsBuf,
|
||||
__ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf,
|
||||
__ubuf__ uint32_t* idx1Buf, uint16_t vfLoop, bool init)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
// 计算直方图0-127 128-255
|
||||
MicroAPI::RegTensor<uint16_t> cout0;
|
||||
MicroAPI::RegTensor<uint16_t> cout1;
|
||||
MicroAPI::Duplicate(cout0, 0);
|
||||
MicroAPI::Duplicate(cout1, 0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Odd;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Odd;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx0;
|
||||
MicroAPI::RegTensor<uint32_t> idx1;
|
||||
// 0x000000fc -> 0xfcfcfcfc
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx0, idx0Buf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx1, idx1Buf);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vreg0U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg1U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg2U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg3U16;
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> vreg0;
|
||||
MicroAPI::RegTensor<uint8_t> vreg1;
|
||||
MicroAPI::RegTensor<uint8_t> vreg2;
|
||||
MicroAPI::RegTensor<uint8_t> vreg3;
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
for (uint16_t i = 0; i < vfLoop; ++i) {
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg1U16,
|
||||
vreg0U16, inputBuf + i * 256);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg3U16,
|
||||
vreg2U16, inputBuf + (i * 256) + 128);
|
||||
|
||||
MicroAPI::DeInterleave(vreg1, vreg0, (MicroAPI::RegTensor<uint8_t>&)vreg0U16,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg2U16);
|
||||
MicroAPI::DeInterleave(vreg3, vreg2, (MicroAPI::RegTensor<uint8_t>&)vreg1U16,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg3U16);
|
||||
|
||||
MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ0, vreg0, (MicroAPI::RegTensor<uint8_t>&)idx0, pregB8);
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ1, vreg1, (MicroAPI::RegTensor<uint8_t>&)idx1, pregB8);
|
||||
|
||||
MicroAPI::MaskReg pregEQ = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::And(pregEQ, pregEQ0, pregEQ1, pregB8);
|
||||
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout0, vreg2, pregEQ);
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout1, vreg2, pregEQ);
|
||||
}
|
||||
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout0U32Even, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout0U32Odd, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout1U32Even, cout1, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout1U32Odd, cout1, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf,
|
||||
cout0U32Even, cout0U32Odd, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf + 128,
|
||||
cout1U32Even, cout1U32Odd, pregB32);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindThirdTargetBinVFImpl(__ubuf__ uint32_t* idx2Buf,
|
||||
__ubuf__ uint32_t* nkValueBuf, __ubuf__ uint32_t* kValue,
|
||||
__ubuf__ uint32_t* histogramsBuf)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx2;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> btmK2;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(btmK2, kValue);
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(4); ++i) {
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::RegTensor<uint32_t> cout;
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdx2;
|
||||
|
||||
MicroAPI::MaskReg pregGE = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::Arange(idxC, i * 64);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(cout, histogramsBuf + i * 64);
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GE>(pregGE, cout, btmK2, pregB32);
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(
|
||||
sqzIdx2,
|
||||
(MicroAPI::RegTensor<uint32_t>&)idxC,
|
||||
pregGE);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(idx2Buf, sqzIdx2, alignIdx2);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(idx2Buf, alignIdx2);
|
||||
|
||||
MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx2;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx2, idx2Buf);
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> idxAll1;
|
||||
MicroAPI::RegTensor<uint32_t> idxPrev2;
|
||||
MicroAPI::RegTensor<uint32_t> prevBinValue;
|
||||
MicroAPI::Duplicate(idxAll1, 1);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> zeroAll;
|
||||
MicroAPI::Duplicate(zeroAll, 0);
|
||||
|
||||
MicroAPI::MaskReg preg2 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::EQ>(preg2, idx2, zeroAll, pregB32);
|
||||
MicroAPI::Sub(idxPrev2, idx2, (MicroAPI::RegTensor<uint32_t>&)idxAll1, pregB32);
|
||||
MicroAPI::ShiftRights(idxPrev2, idxPrev2, (int16_t)24, pregB32);
|
||||
|
||||
MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev2, pregB32);
|
||||
MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg2);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> nextK;
|
||||
MicroAPI::Sub(nextK, btmK2, prevBinValue, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM>(nkValueBuf, nextK, pregB32);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__simd_vf__ void HistogramsLastVFImpl(__ubuf__ uint32_t* histogramsBuf,
|
||||
__ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* idx0Buf,
|
||||
__ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf,
|
||||
uint16_t vfLoop, bool init)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
// 计算直方图0-127 128-255
|
||||
MicroAPI::RegTensor<uint16_t> cout0;
|
||||
MicroAPI::RegTensor<uint16_t> cout1;
|
||||
MicroAPI::Duplicate(cout0, 0);
|
||||
MicroAPI::Duplicate(cout1, 0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Odd;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Odd;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx0;
|
||||
MicroAPI::RegTensor<uint32_t> idx1;
|
||||
MicroAPI::RegTensor<uint32_t> idx2;
|
||||
// 0x000000fc -> 0xfcfcfcfc
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx0, idx0Buf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx1, idx1Buf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idx2, idx2Buf);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vreg0U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg1U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg2U16;
|
||||
MicroAPI::RegTensor<uint32_t> vreg3U16;
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> vreg0;
|
||||
MicroAPI::RegTensor<uint8_t> vreg1;
|
||||
MicroAPI::RegTensor<uint8_t> vreg2;
|
||||
MicroAPI::RegTensor<uint8_t> vreg3;
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
for (uint16_t i = 0; i < vfLoop; ++i) {
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg1U16, vreg0U16, inputBuf + i * 256);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_DINTLV_B16>(vreg3U16,
|
||||
vreg2U16, inputBuf + (i * 256) + 128);
|
||||
|
||||
MicroAPI::DeInterleave(vreg1, vreg0,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg0U16,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg2U16);
|
||||
MicroAPI::DeInterleave(vreg3, vreg2,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg1U16,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vreg3U16);
|
||||
|
||||
MicroAPI::MaskReg pregEQ0 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregEQ1 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregEQ2 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ0, vreg0, (MicroAPI::RegTensor<uint8_t>&)idx0, pregB8);
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ1, vreg1, (MicroAPI::RegTensor<uint8_t>&)idx1, pregB8);
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ2, vreg2, (MicroAPI::RegTensor<uint8_t>&)idx2, pregB8);
|
||||
|
||||
MicroAPI::MaskReg pregEQ0And1 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregEQAll = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::And(pregEQ0And1, pregEQ0, pregEQ1, pregB8);
|
||||
MicroAPI::And(pregEQAll, pregEQ0And1, pregEQ2, pregB8);
|
||||
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout0, vreg3, pregEQAll);
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout1, vreg3, pregEQAll);
|
||||
}
|
||||
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout0U32Even, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout0U32Odd, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout1U32Even, cout1, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout1U32Odd, cout1, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf,
|
||||
cout0U32Even, cout0U32Odd, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf + 128,
|
||||
cout1U32Even, cout1U32Odd, pregB32);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue,
|
||||
__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idx0Buf,
|
||||
__ubuf__ uint32_t* idx1Buf, __ubuf__ uint32_t* idx2Buf,
|
||||
__ubuf__ uint32_t* idx3Buf)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx3;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> btmK3;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(btmK3, kValue);
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(4); ++i) {
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::RegTensor<uint32_t> cout;
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdx3;
|
||||
|
||||
MicroAPI::MaskReg pregGE = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::Arange(idxC, i * 64);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(cout, histogramsBuf + i * 64);
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GE>(pregGE, cout, btmK3, pregB32);
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdx3,
|
||||
(MicroAPI::RegTensor<uint32_t>&)idxC, pregGE);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(idx3Buf, sqzIdx3, alignIdx3);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(idx3Buf, alignIdx3);
|
||||
|
||||
MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idx0;
|
||||
MicroAPI::RegTensor<uint32_t> idx1;
|
||||
MicroAPI::RegTensor<uint32_t> idx2;
|
||||
MicroAPI::RegTensor<uint32_t> idx3;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B32>(idx0, idx0Buf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B32>(idx1, idx1Buf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B32>(idx2, idx2Buf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B32>(idx3, idx3Buf);
|
||||
|
||||
MicroAPI::ShiftLefts(idx0, idx0, (int16_t)24, pregB32);
|
||||
MicroAPI::ShiftLefts(idx1, idx1, (int16_t)16, pregB32);
|
||||
MicroAPI::ShiftLefts(idx2, idx2, (int16_t)8, pregB32);
|
||||
|
||||
// ADD
|
||||
MicroAPI::Add(idx0, idx0, idx1, pregB32);
|
||||
MicroAPI::Add(idx0, idx0, idx2, pregB32);
|
||||
MicroAPI::Add(idx0, idx0, idx3, pregB32);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM>(kValue, idx0, pregB32);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf,
|
||||
__ubuf__ uint32_t* inputBuf, uint32_t beginIdx,
|
||||
__ubuf__ uint32_t* kValue, uint16_t vfLoop)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> kthValue;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(kthValue, kValue);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vregInput;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) {
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::Arange(idxC, beginIdx + i * 64);
|
||||
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(vregInput, inputBuf + i * 64);
|
||||
|
||||
MicroAPI::MaskReg poutGT = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdxOut;
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GT>(poutGT, vregInput, kthValue, pregB32);
|
||||
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdxOut,
|
||||
(MicroAPI::RegTensor<uint32_t>&)idxC, poutGT);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputIdxBuf, sqzIdxOut, alignIdx);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint32_t* outputIdxBuf,
|
||||
__ubuf__ uint32_t* inputBuf, uint32_t beginIdx,
|
||||
__ubuf__ uint32_t* kValue)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> kthValue;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(kthValue, kValue);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vregInput;
|
||||
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::Arange(idxC, beginIdx);
|
||||
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(vregInput, inputBuf);
|
||||
|
||||
MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdxOut;
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::EQ>(poutEQ, vregInput, kthValue, pregB32);
|
||||
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdxOut,
|
||||
(MicroAPI::RegTensor<uint32_t>&)idxC, poutEQ);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputIdxBuf, sqzIdxOut, alignIdx);
|
||||
MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindValueGTOutputVFImpl(__ubuf__ uint32_t* outputValueBuf,
|
||||
__ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue,
|
||||
uint16_t vfLoop)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignValue;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> kthValue;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(kthValue, kValue);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vregInput;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) {
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(vregInput, inputBuf + i * 64);
|
||||
|
||||
MicroAPI::MaskReg poutGT = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> sqzValueOut;
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GT>(poutGT, vregInput, kthValue, pregB32);
|
||||
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzValueOut, vregInput, poutGT);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputValueBuf,
|
||||
sqzValueOut, alignValue);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindValueEQOutputVFImpl(__ubuf__ uint32_t* outputValueBuf,
|
||||
__ubuf__ uint32_t* inputBuf, __ubuf__ uint32_t* kValue)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignValue;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> kthValue;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(kthValue, kValue);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> vregInput;
|
||||
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(vregInput, inputBuf);
|
||||
|
||||
MicroAPI::MaskReg poutEQ = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> sqzValueOut;
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::EQ>(poutEQ, vregInput, kthValue, pregB32);
|
||||
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzValueOut, vregInput, poutEQ);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputValueBuf, sqzValueOut, alignValue);
|
||||
MicroAPI::StoreUnAlignPost(outputValueBuf, alignValue);
|
||||
}
|
||||
|
||||
__aicore__ inline void LiTopKVF(const LocalTensor<uint32_t>& outputIdxLocal,
|
||||
const LocalTensor<uint32_t>& outputValueLocal,
|
||||
const LocalTensor<uint32_t>& inputLocal,
|
||||
const LocalTensor<uint32_t>& tmpIdxLocal,
|
||||
const LocalTensor<uint32_t>& tmpValueLocal,
|
||||
const LocalTensor<uint32_t>& histogramsLocal,
|
||||
const LocalTensor<uint32_t>& idx0Local,
|
||||
const LocalTensor<uint32_t>& idx1Local,
|
||||
const LocalTensor<uint32_t>& idx2Local,
|
||||
const LocalTensor<uint32_t>& idx3Local,
|
||||
const LocalTensor<uint32_t>& nkValueLocal,
|
||||
uint32_t topK,
|
||||
uint32_t s2SeqLen)
|
||||
{
|
||||
__ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* outputValueBuf = (__ubuf__ uint32_t*)outputValueLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* inputBuf = (__ubuf__ uint32_t*)inputLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* tmpIdxBuf = (__ubuf__ uint32_t*)tmpIdxLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* tmpValueBuf = (__ubuf__ uint32_t*)tmpValueLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* idx0Buf = (__ubuf__ uint32_t*)idx0Local.GetPhyAddr();
|
||||
__ubuf__ uint32_t* idx1Buf = (__ubuf__ uint32_t*)idx1Local.GetPhyAddr();
|
||||
__ubuf__ uint32_t* idx2Buf = (__ubuf__ uint32_t*)idx2Local.GetPhyAddr();
|
||||
__ubuf__ uint32_t* idx3Buf = (__ubuf__ uint32_t*)idx3Local.GetPhyAddr();
|
||||
__ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr();
|
||||
|
||||
uint32_t bottomK = s2SeqLen - topK + 1;
|
||||
uint32_t beginIdx = 0;
|
||||
bool flag = true;
|
||||
|
||||
const uint16_t repeatSize8 = 256;
|
||||
const uint16_t repeatSize32 = 64;
|
||||
|
||||
uint16_t histogramsLoopNum = (s2SeqLen + repeatSize8 - 1) / repeatSize8;
|
||||
uint16_t inputLoopNum = (s2SeqLen + repeatSize32 - 1) / repeatSize32;
|
||||
uint16_t topkLoopNum = (topK + 64 - 1) / 64;
|
||||
|
||||
// find kth-value
|
||||
HistogramsFirstVFImpl<uint32_t>(histogramsBuf, inputBuf, histogramsLoopNum, flag);
|
||||
FindFirstTargetBinVFImpl(idx0Buf, nkValueBuf, histogramsBuf, bottomK);
|
||||
HistogramsSecondVFImpl<uint32_t>(histogramsBuf, inputBuf, idx0Buf, histogramsLoopNum, flag);
|
||||
FindSecondTargetBinVFImpl(idx1Buf, nkValueBuf, nkValueBuf, histogramsBuf);
|
||||
HistogramsThirdVFImpl<uint32_t>(histogramsBuf, inputBuf, idx0Buf, idx1Buf, histogramsLoopNum, flag);
|
||||
FindThirdTargetBinVFImpl(idx2Buf, nkValueBuf, nkValueBuf, histogramsBuf);
|
||||
HistogramsLastVFImpl<uint32_t>(histogramsBuf, inputBuf, idx0Buf, idx1Buf, idx2Buf, histogramsLoopNum, flag);
|
||||
FindKthVFImpl(nkValueBuf, histogramsBuf, idx0Buf, idx1Buf, idx2Buf, idx3Buf);
|
||||
|
||||
// filter
|
||||
// 输出大于k-value的值value
|
||||
FindValueGTOutputVFImpl(outputValueBuf, inputBuf, nkValueBuf, inputLoopNum);
|
||||
// value-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数
|
||||
int64_t arValueNum = AscendC::GetSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
// value-剩余需要输出等于k-value的数量
|
||||
int64_t remainValueNum = topK - (arValueNum / sizeof(uint32_t));
|
||||
for (uint16_t i = 0; i < inputLoopNum; ++i) {
|
||||
int64_t arValueNumPerLoop = AscendC::GetSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
if (((arValueNumPerLoop - arValueNum) / sizeof(uint32_t)) < remainValueNum) {
|
||||
// 调用一次查找等于k-value情况的过程
|
||||
FindValueEQOutputVFImpl(outputValueBuf, inputBuf + i * 64, nkValueBuf);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 输出大于k-value的值idx
|
||||
FindIdxGTOutputVFImpl(outputIdxBuf, inputBuf, (uint32_t)(0), nkValueBuf, inputLoopNum);
|
||||
// idx-当前偏移大于k-value的值在AR特殊寄存器中的有效字节数
|
||||
int64_t arIdxNum = AscendC::GetSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
int64_t remainIdxNum = topK - (arIdxNum / sizeof(uint32_t));
|
||||
for (uint16_t i = 0; i < inputLoopNum; ++i) {
|
||||
int64_t arIdxNumPerLoop = AscendC::GetSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
if (((arIdxNumPerLoop - arIdxNum) / sizeof(uint32_t)) < remainIdxNum) {
|
||||
// 调用一次查找等于k-value情况的过程
|
||||
beginIdx = i * 64;
|
||||
FindIdxEQOutputVFImpl(outputIdxBuf, inputBuf + i * 64, beginIdx, nkValueBuf);
|
||||
} else {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,465 @@
|
||||
/**
|
||||
* 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 vf_top_k_16_gather.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef VF_TOP_K_16_GATHER_H
|
||||
#define VF_TOP_K_16_GATHER_H
|
||||
|
||||
namespace topkb16gather {
|
||||
|
||||
template<typename T>
|
||||
__simd_vf__ void HistogramsHighVFImpl(__ubuf__ uint32_t* histogramsBuf,
|
||||
__ubuf__ uint16_t* inputBuf,
|
||||
uint16_t vfLoop, bool init)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
// 计算直方图cout0 0-127 cout1 128-255
|
||||
MicroAPI::RegTensor<uint16_t> cout0;
|
||||
MicroAPI::RegTensor<uint16_t> cout1;
|
||||
MicroAPI::Duplicate(cout0, 0);
|
||||
MicroAPI::Duplicate(cout1, 0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Odd;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Odd;
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> vregHigh;
|
||||
MicroAPI::RegTensor<uint16_t> vregLow;
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
for (uint16_t i = 0; i < vfLoop; ++i) {
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_DINTLV_B8>(vregLow, vregHigh, inputBuf + i * 256);
|
||||
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout0,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vregHigh,
|
||||
pregB8);
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout1,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vregHigh,
|
||||
pregB8);
|
||||
}
|
||||
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout0U32Even, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout0U32Odd, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout1U32Even, cout1, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout1U32Odd, cout1, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf,
|
||||
cout0U32Even, cout0U32Odd, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf + 128,
|
||||
cout1U32Even, cout1U32Odd, pregB32);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindHighTargetBinVFImpl(__ubuf__ uint32_t* idxHighBuf,
|
||||
__ubuf__ uint32_t* nkValueBuf,
|
||||
__ubuf__ uint32_t* histogramsBuf,
|
||||
uint32_t bottomK)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::MaskReg pregGE;
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdxHigh;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> btmK;
|
||||
MicroAPI::Duplicate(btmK, bottomK);
|
||||
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::RegTensor<uint32_t> cout;
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdxHigh;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(4); ++i) {
|
||||
MicroAPI::Arange(idxC, i * 64);
|
||||
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(cout, histogramsBuf + i * 64);
|
||||
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GE>(pregGE, cout, btmK, pregB32);
|
||||
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(
|
||||
sqzIdxHigh, (MicroAPI::RegTensor<uint32_t>&)idxC, pregGE);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(idxHighBuf, sqzIdxHigh, alignIdxHigh);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(idxHighBuf, alignIdxHigh);
|
||||
|
||||
MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idxHigh;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idxHigh, idxHighBuf);
|
||||
|
||||
MicroAPI::RegTensor<uint8_t> idxAll1;
|
||||
MicroAPI::RegTensor<uint32_t> idxPrev0;
|
||||
MicroAPI::RegTensor<uint32_t> prevBinValue;
|
||||
MicroAPI::Duplicate(idxAll1, 1);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> zeroAll;
|
||||
MicroAPI::Duplicate(zeroAll, 0);
|
||||
|
||||
MicroAPI::MaskReg preg0 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::EQ>(preg0, idxHigh, zeroAll, pregB32);
|
||||
MicroAPI::Sub(idxPrev0, idxHigh, (MicroAPI::RegTensor<uint32_t>&)idxAll1, pregB32);
|
||||
MicroAPI::ShiftRights(idxPrev0, idxPrev0, (int16_t)24, pregB32);
|
||||
|
||||
MicroAPI::Gather(prevBinValue, histogramsBuf, idxPrev0, pregB32);
|
||||
MicroAPI::Select(prevBinValue, zeroAll, prevBinValue, preg0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> nextK;
|
||||
MicroAPI::Sub(nextK, btmK, prevBinValue, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM>(nkValueBuf, nextK, pregB32);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__simd_vf__ void HistogramsLowVFImpl(__ubuf__ uint32_t* histogramsBuf,
|
||||
__ubuf__ uint16_t* inputBuf, __ubuf__ uint32_t* idxHighBuf,
|
||||
uint16_t vfLoop, bool init)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB8 = MicroAPI::CreateMask<uint8_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::MaskReg pregEQ;
|
||||
|
||||
// 计算直方图0-127 128-255
|
||||
MicroAPI::RegTensor<uint16_t> cout0;
|
||||
MicroAPI::RegTensor<uint16_t> cout1;
|
||||
MicroAPI::Duplicate(cout0, 0);
|
||||
MicroAPI::Duplicate(cout1, 0);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout0U32Odd;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Even;
|
||||
MicroAPI::RegTensor<uint32_t> cout1U32Odd;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idxHigh;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idxHigh, idxHighBuf);
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> vregHigh;
|
||||
MicroAPI::RegTensor<uint16_t> vregLow;
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_EVEN = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
static constexpr MicroAPI::CastTrait CAST_TRAIT_UINT16_TOUINT32_ODD = {MicroAPI::RegLayout::ONE,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
for (uint16_t i = 0; i < vfLoop; ++i) {
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_DINTLV_B8>(vregLow, vregHigh, inputBuf + i * 256);
|
||||
|
||||
MicroAPI::Compare<uint8_t, CMPMODE::EQ>(pregEQ,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vregHigh,
|
||||
(MicroAPI::RegTensor<uint8_t>&)idxHigh, pregB8);
|
||||
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN0,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout0,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vregLow, pregEQ);
|
||||
MicroAPI::Histograms<uint8_t, uint16_t, MicroAPI::HistogramsBinType::BIN1,
|
||||
MicroAPI::HistogramsType::ACCUMULATE>(cout1,
|
||||
(MicroAPI::RegTensor<uint8_t>&)vregLow, pregEQ);
|
||||
}
|
||||
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout0U32Even, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout0U32Odd, cout0, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_EVEN>(cout1U32Even, cout1, pregB16);
|
||||
MicroAPI::Cast<uint32_t, uint16_t, CAST_TRAIT_UINT16_TOUINT32_ODD>(cout1U32Odd, cout1, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf,
|
||||
cout0U32Even, cout0U32Odd, pregB32);
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_INTLV_B32>(histogramsBuf + 128,
|
||||
cout1U32Even, cout1U32Odd, pregB32);
|
||||
}
|
||||
|
||||
__simd_vf__ void FindKthVFImpl(__ubuf__ uint32_t* kValue,
|
||||
__ubuf__ uint32_t* histogramsBuf, __ubuf__ uint32_t* idxHighBuf,
|
||||
__ubuf__ uint32_t* idxLowBuf)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::MaskReg pregGE;
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdxLow;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> btmK;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(btmK, kValue);
|
||||
|
||||
MicroAPI::RegTensor<int32_t> idxC;
|
||||
MicroAPI::RegTensor<uint32_t> cout;
|
||||
MicroAPI::RegTensor<uint32_t> sqzIdxLow;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(4); ++i) {
|
||||
MicroAPI::Arange(idxC, i * 64);
|
||||
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_NORM>(cout, histogramsBuf + i * 64);
|
||||
|
||||
MicroAPI::Compare<uint32_t, CMPMODE::GE>(pregGE, cout, btmK, pregB32);
|
||||
|
||||
MicroAPI::Squeeze<uint32_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdxLow,
|
||||
(MicroAPI::RegTensor<uint32_t>&)idxC, pregGE);
|
||||
MicroAPI::StoreUnAlign<uint32_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(idxLowBuf, sqzIdxLow, alignIdxLow);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(idxLowBuf, alignIdxLow);
|
||||
|
||||
MicroAPI::LocalMemBar<AscendC::MicroAPI::MemType::VEC_STORE, AscendC::MicroAPI::MemType::VEC_LOAD>();
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idxHigh;
|
||||
MicroAPI::RegTensor<uint32_t> idxLow;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B8>(idxHigh, idxHighBuf);
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B16>(idxLow, idxLowBuf);
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> idxTmp;
|
||||
MicroAPI::Duplicate(idxTmp, 0xff00);
|
||||
|
||||
MicroAPI::And(idxHigh, idxHigh, (MicroAPI::RegTensor<uint32_t>&)idxTmp, pregB32);
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> idxK;
|
||||
MicroAPI::Add(idxK, idxHigh, idxLow, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM_B16>(kValue, idxK, pregB32);
|
||||
}
|
||||
|
||||
/**
|
||||
输出所有大于的kth-value的Index
|
||||
*/
|
||||
__simd_vf__ void FindIdxGTOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf,
|
||||
__ubuf__ uint16_t* inputValueBuf, uint16_t beginIdx,
|
||||
__ubuf__ uint32_t* kValue, uint16_t vfLoop)
|
||||
{
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::MaskReg poutGT;
|
||||
|
||||
MicroAPI::ClearSpr<AscendC::SpecialPurposeReg::AR>();
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> kthValue;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B16>(kthValue, kValue);
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> vregInput;
|
||||
MicroAPI::RegTensor<int16_t> idxC;
|
||||
MicroAPI::RegTensor<uint16_t> sqzIdxOut;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) {
|
||||
MicroAPI::Arange(idxC, beginIdx + i * 128);
|
||||
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_NORM>(vregInput, inputValueBuf + i * 128);
|
||||
|
||||
MicroAPI::Compare<uint16_t, CMPMODE::GT>(poutGT, vregInput, (MicroAPI::RegTensor<uint16_t>&)kthValue, pregB16);
|
||||
|
||||
MicroAPI::Squeeze<uint16_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdxOut,
|
||||
(MicroAPI::RegTensor<uint16_t>&)idxC, poutGT);
|
||||
MicroAPI::StoreUnAlign<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputIdxBuf, sqzIdxOut, alignIdx);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx);
|
||||
}
|
||||
|
||||
/**
|
||||
输出所有等于的kth-value的Index
|
||||
*/
|
||||
__simd_vf__ void FindIdxEQOutputVFImpl(__ubuf__ uint16_t* outputIdxBuf,
|
||||
__ubuf__ uint16_t* inputValueBuf, uint16_t beginIdx,
|
||||
__ubuf__ uint32_t* kValue, uint16_t vfLoop)
|
||||
{
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::MaskReg poutEQ;
|
||||
|
||||
MicroAPI::UnalignRegForStore alignIdx;
|
||||
|
||||
MicroAPI::RegTensor<uint32_t> kthValue;
|
||||
MicroAPI::LoadAlign<uint32_t, MicroAPI::LoadDist::DIST_BRC_B16>(kthValue, kValue);
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> vregInput;
|
||||
MicroAPI::RegTensor<int16_t> idxC;
|
||||
MicroAPI::RegTensor<uint16_t> sqzIdxOut;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) {
|
||||
MicroAPI::Arange(idxC, beginIdx + i * 128);
|
||||
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_NORM>(vregInput, inputValueBuf + i * 128);
|
||||
|
||||
MicroAPI::Compare<uint16_t, CMPMODE::EQ>(poutEQ, vregInput, (MicroAPI::RegTensor<uint16_t>&)kthValue, pregB16);
|
||||
|
||||
MicroAPI::Squeeze<uint16_t, MicroAPI::GatherMaskMode::STORE_REG>(sqzIdxOut,
|
||||
(MicroAPI::RegTensor<uint16_t>&)idxC, poutEQ);
|
||||
MicroAPI::StoreUnAlign<uint16_t, MicroAPI::PostLiteral::POST_MODE_UPDATE>(outputIdxBuf, sqzIdxOut, alignIdx);
|
||||
}
|
||||
MicroAPI::StoreUnAlignPost(outputIdxBuf, alignIdx);
|
||||
}
|
||||
|
||||
/**
|
||||
输出最终的Value
|
||||
*/
|
||||
__simd_vf__ void FindValueOutputVFImpl(__ubuf__ uint16_t* outputValueBuf,
|
||||
__ubuf__ uint16_t* inputValueBuf,
|
||||
__ubuf__ uint16_t* tmpIdxBuf, uint16_t vfLoop)
|
||||
{
|
||||
MicroAPI::MaskReg pregB16 = MicroAPI::CreateMask<uint16_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> tmpIdx;
|
||||
MicroAPI::RegTensor<uint16_t> outputValue;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) {
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_NORM>(tmpIdx, tmpIdxBuf + i * 128);
|
||||
|
||||
MicroAPI::Gather(outputValue, inputValueBuf, tmpIdx, pregB16);
|
||||
|
||||
MicroAPI::StoreAlign<uint16_t, MicroAPI::StoreDist::DIST_NORM>(outputValueBuf + i * 128, outputValue, pregB16);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
输出最终的Idx
|
||||
*/
|
||||
__simd_vf__ void FindRealIndexVFImpl(__ubuf__ uint32_t* outputIdxBuf,
|
||||
__ubuf__ uint16_t* tmpIdxBuf, __ubuf__ uint32_t* hisIdxBuf,
|
||||
uint32_t topK, uint32_t loopIndex, uint16_t vfLoop)
|
||||
{
|
||||
MicroAPI::MaskReg pregB32 = MicroAPI::CreateMask<uint32_t, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::MaskReg pregNow;
|
||||
MicroAPI::MaskReg pregHis;
|
||||
|
||||
MicroAPI::RegTensor<uint16_t> tmpIdx;
|
||||
MicroAPI::RegTensor<uint32_t> outputGatherIdx;
|
||||
MicroAPI::RegTensor<uint32_t> outputAddsIdx;
|
||||
|
||||
for (uint16_t i = 0; i < (uint16_t)(vfLoop); ++i) {
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>(tmpIdx, tmpIdxBuf + i * 64);
|
||||
|
||||
MicroAPI::Compares<uint32_t, CMPMODE::GT>(pregNow, (MicroAPI::RegTensor<uint32_t>&)tmpIdx, topK - 1, pregB32);
|
||||
MicroAPI::Xor(pregHis, pregNow, pregB32, pregB32);
|
||||
|
||||
MicroAPI::Gather(outputGatherIdx, hisIdxBuf, (MicroAPI::RegTensor<uint32_t>&)tmpIdx, pregHis);
|
||||
MicroAPI::Adds(outputAddsIdx, (MicroAPI::RegTensor<uint32_t>&)tmpIdx, loopIndex, pregNow);
|
||||
|
||||
MicroAPI::Add(outputGatherIdx, outputGatherIdx, outputAddsIdx, pregB32);
|
||||
|
||||
MicroAPI::StoreAlign<uint32_t, MicroAPI::StoreDist::DIST_NORM>(outputIdxBuf + i * 64, outputGatherIdx, pregB32);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief LiTopKVF 对一个validLen的输入进行topk算法，输出idx_tmp
|
||||
* @param tmpIdxLocal Temp阶段输出的TopKIndex;如果s2SeqLen < 16K作为最终输出 validLen * 2B
|
||||
* @param outputValueLocal 如果s2SeqLen > 16K并且是首轮输出Value topK * 2B
|
||||
* @param inputValueLocal 输入Value validLen * 2B
|
||||
* @param histogramsLocal 直方图 256 * 4B
|
||||
* @param idxHighLocal 目标桶高八位 256 * 4B
|
||||
* @param idxLowLocal 目标桶低八位 256 * 4B
|
||||
* @param nkValueLocal 存储next_k的值 64 * 4B
|
||||
* @param topK topK元素
|
||||
* @param validLen 有效元素个数:LICommon::Align(topkCountAlign256_ + validTrunkLen, (uint32_t)256)
|
||||
*/
|
||||
template<bool ISOUTVALUE> // 是否输出VALUE
|
||||
__aicore__ inline void LiTopKVF(const LocalTensor<uint16_t>& tmpIdxLocal,
|
||||
const LocalTensor<uint16_t>& outputValueLocal,
|
||||
const LocalTensor<uint16_t>& inputValueLocal,
|
||||
const LocalTensor<uint32_t>& histogramsLocal,
|
||||
const LocalTensor<uint32_t>& idxHighLocal,
|
||||
const LocalTensor<uint32_t>& idxLowLocal,
|
||||
const LocalTensor<uint32_t>& nkValueLocal,
|
||||
uint32_t topK,
|
||||
uint32_t validLen)
|
||||
{
|
||||
__ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr();
|
||||
__ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr();
|
||||
__ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* histogramsBuf = (__ubuf__ uint32_t*)histogramsLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* idxHighBuf = (__ubuf__ uint32_t*)idxHighLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* idxLowBuf = (__ubuf__ uint32_t*)idxLowLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* nkValueBuf = (__ubuf__ uint32_t*)nkValueLocal.GetPhyAddr();
|
||||
|
||||
uint32_t bottomK = validLen - topK + 1;
|
||||
uint32_t beginIdx = 0;
|
||||
bool flag = true;
|
||||
|
||||
const uint16_t repeatSize8 = 256;
|
||||
const uint16_t repeatSize16 = 128;
|
||||
const uint16_t repeatSize32 = 64;
|
||||
|
||||
uint16_t histogramsLoopNum = (validLen + repeatSize8 - 1) / repeatSize8;
|
||||
uint16_t inputLoopNum = (validLen + repeatSize16 - 1) / repeatSize16;
|
||||
uint16_t topkLoopNum = (topK + repeatSize32 - 1) / repeatSize32;
|
||||
uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16;
|
||||
|
||||
// find kth-value
|
||||
HistogramsHighVFImpl<uint16_t>(histogramsBuf, inputValueBuf, histogramsLoopNum, flag);
|
||||
FindHighTargetBinVFImpl(idxHighBuf, nkValueBuf, histogramsBuf, bottomK);
|
||||
|
||||
HistogramsLowVFImpl<uint16_t>(histogramsBuf, inputValueBuf, idxHighBuf, histogramsLoopNum, flag);
|
||||
FindKthVFImpl(nkValueBuf, histogramsBuf, idxHighBuf, idxLowBuf);
|
||||
|
||||
// filter
|
||||
int32_t count = LICommon::Align(topK, (uint32_t)128) - topK / 128 * 128;
|
||||
AscendC::Duplicate(tmpIdxLocal[topK / 128 * 128], (uint16_t)(0), count);
|
||||
// 输出大于k-value的值idx
|
||||
FindIdxGTOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum);
|
||||
// 输出等于k-value的值idx
|
||||
FindIdxEQOutputVFImpl(tmpIdxBuf, inputValueBuf, (uint32_t)(0), nkValueBuf, inputLoopNum);
|
||||
|
||||
// 是否输出Value
|
||||
if constexpr (ISOUTVALUE) {
|
||||
FindValueOutputVFImpl(outputValueBuf, inputValueBuf, tmpIdxBuf, topkLoopNum16);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief 通过idx_tmp gather出实际的TopKIndex，s2SeqLen > 16K才会执行
|
||||
* @param outputIdxLocal 输出Idx 有效:topK * 2B
|
||||
* @param outputValueLocal 输出Value topK * 2B(以后需要输出实际value使用)
|
||||
* @param inputValueLocal 输入Value validLen * 2B
|
||||
* @param tmpIdxLocal 本轮tmpIdx输入 validLen * 2B (0 ~ validLen - 1)
|
||||
* @param hisIdxLocal 上一轮实际Idx输入 有效:topK * 4B
|
||||
* @param topK topK元素个数
|
||||
* @param loopBasicIdx 当前循环需要加上得基准Index
|
||||
* @param validLen 有效元素个数
|
||||
*/
|
||||
__aicore__ inline void LiTopKGatherVF(const LocalTensor<uint32_t>& outputIdxLocal,
|
||||
const LocalTensor<uint16_t>& outputValueLocal,
|
||||
const LocalTensor<uint16_t>& inputValueLocal,
|
||||
const LocalTensor<uint16_t>& tmpIdxLocal,
|
||||
const LocalTensor<uint32_t>& hisIdxLocal,
|
||||
uint32_t topK,
|
||||
uint32_t loopBasicIdx,
|
||||
uint32_t validLen)
|
||||
{
|
||||
__ubuf__ uint32_t* outputIdxBuf = (__ubuf__ uint32_t*)outputIdxLocal.GetPhyAddr();
|
||||
__ubuf__ uint16_t* outputValueBuf = (__ubuf__ uint16_t*)outputValueLocal.GetPhyAddr();
|
||||
__ubuf__ uint16_t* inputValueBuf = (__ubuf__ uint16_t*)inputValueLocal.GetPhyAddr();
|
||||
__ubuf__ uint16_t* tmpIdxBuf = (__ubuf__ uint16_t*)tmpIdxLocal.GetPhyAddr();
|
||||
__ubuf__ uint32_t* hisIdxBuf = (__ubuf__ uint32_t*)hisIdxLocal.GetPhyAddr();
|
||||
|
||||
const uint16_t repeatSize32 = 64;
|
||||
const uint16_t repeatSize16 = 128;
|
||||
uint16_t topkLoopNum16 = (topK + repeatSize16 - 1) / repeatSize16;
|
||||
uint16_t topkLoopNum32 = (topK + repeatSize32 - 1) / repeatSize32;
|
||||
|
||||
FindRealIndexVFImpl(outputIdxBuf, tmpIdxBuf, hisIdxBuf, topK, loopBasicIdx, topkLoopNum32);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
Reference in New Issue
Block a user