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

init v0.23.0

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

181
csrc/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_topk.h Normal file
View File

@@ -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

830
csrc/attention/lightning_indexer/op_kernel/arch35/vf/lightning_indexer_vector1.h Normal file
View File

@@ -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> (&regArray)[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
View 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

465
csrc/attention/lightning_indexer/op_kernel/arch35/vf/vf_topk_16_gather.h Normal file
View File

@@ -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