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

29
csrc/common/include/kernel/common.h Normal file
View File

@@ -0,0 +1,29 @@
/**
 * Copyright (c) 2025 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 common.h
* \brief
*/
#ifndef INCLUDE_COMMON_H
#define INCLUDE_COMMON_H
#define CONST_2 2
#define SET_FLAG(trigger, waiter, e) AscendC::SetFlag<AscendC::HardEvent::trigger##_##waiter>((e))
#define WAIT_FLAG(trigger, waiter, e) AscendC::WaitFlag<AscendC::HardEvent::trigger##_##waiter>((e))
#define PIPE_BARRIER(pipe) AscendC::PipeBarrier<PIPE_##pipe>()
#ifndef FORCE_INLINE
#define FORCE_INLINE inline __attribute__((always_inline))
#endif
#endif

117
csrc/common/include/kernel/common_func.h Normal file
View File

@@ -0,0 +1,117 @@
/**
 * Copyright (c) 2025 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 common_func.h
* \brief
*/
#ifndef INCLUDE_COMMON_FUNC_H
#define INCLUDE_COMMON_FUNC_H
#include <limits>
#include <type_traits>
#ifdef __CCE_KT_TEST__
#include "stub_def.h"
#include "stub_fun.h"
#else
#include "kernel_macros.h"
#endif
template <uint32_t ALIGN, typename T = uint32_t>
inline __aicore__ T RoundUp(const T val)
{
static_assert(ALIGN != 0, "align must not be zero");
static_assert(std::is_arithmetic<T>::value, "T must be an arithmetic type");
T align = ALIGN;
if (val + align - 1 < val) {
return val;
}
return (val + align - 1) / align * align;
}
template <typename T>
inline __aicore__ T RoundUp(const T val, const T align)
{
static_assert(std::is_arithmetic<T>::value, "T must be an arithmetic type");
if (align == 0 || val + align - 1 < val) {
return val;
}
return (val + align - 1) / align * align;
}
template <uint32_t DIVISOR, typename T = uint32_t>
inline __aicore__ T CeilDiv(const T dividend)
{
static_assert(DIVISOR != 0, "align must not be zero");
static_assert(std::is_arithmetic<T>::value, "T must be an arithmetic type");
T divisor = DIVISOR;
if (dividend + divisor - 1 < dividend) {
return dividend;
}
return (dividend + divisor - 1) / divisor;
}
template <typename T>
constexpr T T_MAX = std::numeric_limits<T>::max();
template <typename T>
inline __aicore__ T CeilDiv(const T dividend, const T divisor)
{
static_assert(std::is_arithmetic<T>::value, "T must be an arithmetic type");
if (divisor == 0 || dividend + divisor - 1 < dividend) {
return T_MAX<T>;
}
return (dividend + divisor - 1) / divisor;
}
template <typename T>
__aicore__ inline T Min(const T lhs, const T rhs)
{
return lhs < rhs ? lhs : rhs;
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint32_t BlockSize()
{
return 32 / sizeof(Dtype);
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint32_t MatrixSize()
{
return 512 / sizeof(Dtype);
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint64_t BlockSizeRoundUp(uint64_t num)
{
return (num + BlockSize<Dtype>() - 1) / BlockSize<Dtype>() * BlockSize<Dtype>();
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint64_t NumBlocksRoundUp(uint64_t num)
{
return (num + BlockSize<Dtype>() - 1) / BlockSize<Dtype>();
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint64_t MatrixSizeRoundUp(uint64_t num)
{
return (num + MatrixSize<Dtype>() - 1) / MatrixSize<Dtype>() * MatrixSize<Dtype>();
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint64_t NumMatrixsRoundUp(uint64_t num)
{
return (num + MatrixSize<Dtype>() - 1) / MatrixSize<Dtype>();
}
template <typename Dtype> __aicore__ __attribute__((always_inline)) inline uint64_t L0HalfSize()
{
return 32 * 1024 / sizeof(Dtype);
}
#endif

121
csrc/common/include/kernel/dropmask.h Normal file
View File

@@ -0,0 +1,121 @@
/**
 * Copyright (c) 2025 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 dropmask.h
* \brief
*/
#ifndef DROPMASK_H
#define DROPMASK_H
#include "util.h"
using AscendC::DROPOUT_MODE_BIT_MISALIGN;
using AscendC::DropOutShapeInfo;
using AscendC::DropOut;
struct DropMaskInfo {
// for compute dropout mask offset
// 参数按B N G S1 S2全部切分设置进行偏移计算，没有切分的轴对应的参数设置为合适的0或者原始值
int64_t n2G; // n2 * g
int64_t gSize; // g
int64_t s1Size; // s1
int64_t s2Size; // s2
int64_t gOutIdx; // g out index
int64_t bSSOffset; // boidx * s1 * s2 ===bSSOffset
int64_t n2OutIdx; // n out index
int64_t s1OutIdx; // s1 out index ===s1oIdx
int64_t s1InnerIdx; // s1 inner index, 配比 ===loopIdx
int64_t s1BaseSize; // S1基本块大小
int64_t splitS1BaseSize; // s1 split size ===vec1S1BaseSize
int64_t s2StartIdx; // s2 start index
int64_t s2Idx; // s2 index =====s2LoopCount
int64_t s2BaseNratioSize; // s2的配比长度: s2BaseSize(S2基本块大小) * nRatio
// for copy in dropout mask
uint32_t s1CopySize;
uint32_t s2CopySize;
int64_t s2TotalSize;
// for compute dropout mask
uint32_t firstAxis;
uint32_t lstAxis;
uint32_t maskLstAxis;
int64_t vecCoreOffset = 0;
float keepProb;
bool boolMode;
};
template <bool hasDrop>
__aicore__ inline int64_t ComputeDropOffset(DropMaskInfo &dropMaskInfo)
{
if constexpr (hasDrop == true) {
// boidx * n2 * g* s1 * s2
int64_t bOffset = dropMaskInfo.bSSOffset * dropMaskInfo.n2G;
// n2oIdx * g * s1 *s2
int64_t n2Offset = dropMaskInfo.n2OutIdx * dropMaskInfo.gSize * dropMaskInfo.s1Size * dropMaskInfo.s2Size;
// goIdx * s1 * s2
int64_t gOffset = dropMaskInfo.gOutIdx * dropMaskInfo.s1Size * dropMaskInfo.s2Size;
// s1oIdx * s1BaseSize * s2Size + s1innerindex * vec1S1BaseSize * s2Size
int64_t s1Offset = (dropMaskInfo.s1OutIdx * dropMaskInfo.s1BaseSize + dropMaskInfo.vecCoreOffset +
dropMaskInfo.s1InnerIdx * dropMaskInfo.splitS1BaseSize) * dropMaskInfo.s2Size;
// s2StartIdx + s2index * s2BaseNratioSize
int64_t s2Offset = dropMaskInfo.s2StartIdx + dropMaskInfo.s2Idx * dropMaskInfo.s2BaseNratioSize;
return bOffset + n2Offset + gOffset + s1Offset + s2Offset;
} else {
return 0;
}
}
template <bool hasDrop>
__aicore__ inline void CopyInDropMask(LocalTensor<uint8_t>&dstTensor, GlobalTensor<uint8_t>& srcBoolTensor,
GlobalTensor<uint8_t>& srcByteTensor, DropMaskInfo &dropMaskInfo, int64_t alignedSize = blockBytes)
{
if constexpr (hasDrop == true) {
int64_t dropMaskOffset = ComputeDropOffset<hasDrop>(dropMaskInfo);
if (unlikely(dropMaskInfo.boolMode)) {
BoolCopyIn(dstTensor, srcBoolTensor, dropMaskOffset,
dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize);
} else {
Bit2Int8CopyIn(dstTensor, srcByteTensor, dropMaskOffset, 1,
dropMaskInfo.s1CopySize, dropMaskInfo.s2CopySize, dropMaskInfo.s2TotalSize, alignedSize);
}
return;
}
}
template <typename T, bool hasDrop>
__aicore__ inline void ComputeDropMask(LocalTensor<T>& dstTensor, LocalTensor<T>& srcTensor,
LocalTensor<uint8_t>& dropoutBuffer, LocalTensor<uint8_t>& tmpDropBuffer, DropMaskInfo &dropMaskInfo)
{
if constexpr (hasDrop == true) {
DropOutShapeInfo dropOutShapeInfo;
dropOutShapeInfo.firstAxis = dropMaskInfo.firstAxis;
dropOutShapeInfo.srcLastAxis = dropMaskInfo.lstAxis;
if (unlikely(dropMaskInfo.boolMode)) {
dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis, blockBytes) * blockBytes;
DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo);
} else {
dropOutShapeInfo.maskLastAxis = CeilDiv(dropMaskInfo.maskLstAxis / byteBitRatio, blockBytes) * blockBytes;
if (likely(dropMaskInfo.lstAxis / byteBitRatio % blockBytes == 0)) {
DropOut(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer, dropMaskInfo.keepProb, dropOutShapeInfo);
} else {
DropOut<T, false, DROPOUT_MODE_BIT_MISALIGN>(dstTensor, srcTensor, dropoutBuffer, tmpDropBuffer,
dropMaskInfo.keepProb, dropOutShapeInfo);
}
}
return;
}
}
#endif // DROPMASK_H

169
csrc/common/include/kernel/gm_to_l1_iterator.h Normal file
View File

@@ -0,0 +1,169 @@
/**
 * Copyright (c) 2025 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 gm_to_l1_iterator.h
* \brief
*/
#ifndef GM_TO_L1_ITERATOR_H
#define GM_TO_L1_ITERATOR_H
#include "iterator.h"
constexpr uint32_t STRIDE_LIMIT_H = 65536;
// Partial specialization for V220, ND_in, ND_out
template <ArchType ArchTag, typename DataType>
struct gm_to_l1<ArchTag, DataType, DataFormatT::ND, DataFormatT::ND> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
__aicore__ gm_to_l1(AscendC::LocalTensor<DataType> l1Tensor,
AscendC::GlobalTensor<DataType> gmTensor,
uint32_t nTileActual,
uint32_t nTileCeil,
uint32_t nVal,
uint32_t dTileActual,
uint32_t dTileCeil,
uint32_t dVal)
{
AscendC::DataCopy(l1Tensor,
gmTensor,
AscendC::DataCopyParams(1, // nBurst
CeilDiv<BLOCK_SIZE>(nTileActual * dTileActual), // lenBurst
0, // srcGap
0)); // dstGap
};
};
// Partial specialization for NZ_in, NZ_out
template <ArchType ArchTag, typename DataType>
struct gm_to_l1<ArchTag, DataType, DataFormatT::NZ, DataFormatT::NZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
__aicore__ gm_to_l1(AscendC::LocalTensor<DataType> l1Tensor,
AscendC::GlobalTensor<DataType> gmTensor,
uint32_t nTileActual,
uint32_t nTileCeil,
uint32_t nVal,
uint32_t dTileActual,
uint32_t dTileCeil,
uint32_t dVal)
{
uint64_t srcStride = nTileCeil - nTileActual;
if (srcStride < STRIDE_LIMIT_H) {
AscendC::DataCopy(l1Tensor, gmTensor,
AscendC::DataCopyParams(dTileActual / BLOCK_SIZE, // nBurst
nTileActual, // lenBurst
nTileCeil - nTileActual, // srcGap
0)); // dstGap
} else {
for (uint64_t i = 0; i < dTileActual / BLOCK_SIZE; i++) {
uint64_t dstOffset = i * nTileActual * BLOCK_SIZE;
uint64_t srcOffset = i * nTileCeil * BLOCK_SIZE;
AscendC::DataCopy(l1Tensor[dstOffset], gmTensor[srcOffset],
AscendC::DataCopyParams(1, // nBurst
nTileActual, // lenBurst
0, // srcGap
0)); // dstGap
}
}
};
};
// Partial specialization for V220, ND_in, ND_out
template <ArchType ArchTag, typename DataType>
struct gm_to_l1<ArchTag, DataType, DataFormatT::ND, DataFormatT::NZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
__aicore__ gm_to_l1(AscendC::LocalTensor<DataType> l1Tensor,
AscendC::GlobalTensor<DataType> gmTensor,
uint32_t nTileActual,
uint32_t nTileCeil,
uint32_t nVal,
uint32_t dTileActual,
uint32_t dTileCeil,
uint32_t dVal)
{
if (dVal < STRIDE_LIMIT_H) {
AscendC::DataCopy(l1Tensor,
gmTensor,
AscendC::Nd2NzParams(1, // ndNum
nTileActual, // nValue
dTileActual, // dValue
0, // srcNdMatrixStride, unused
dVal, // srcDValue
nTileCeil, // dstNzC0Stride
1, // dstNzNStride
0)); // dstNzMatrixStride, unused
} else {
for (uint32_t i = 0; i < nTileActual; i++) {
AscendC::DataCopy(l1Tensor[i * BLOCK_SIZE],
gmTensor[i * dVal],
AscendC::Nd2NzParams(1, // ndNum
1, // nValue
dTileActual, // dValue
0, // srcNdMatrixStride, unused
0, // srcDValue
nTileCeil, // dstNzC0Stride
0, // dstNzNStride
0)); // dstNzMatrixStride, unused
}
}
};
};
// Partial specialization for V220, ND_in, NZ_out
template <ArchType ArchTag, typename DataType>
struct gm_to_l1<ArchTag, DataType, DataFormatT::ND, DataFormatT::ZN> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
__aicore__ gm_to_l1(AscendC::LocalTensor<DataType> l1Tensor,
AscendC::GlobalTensor<DataType> gmTensor,
uint32_t nTileActual,
uint32_t nTileCeil,
uint32_t nVal,
uint32_t dTileActual,
uint32_t dTileCeil,
uint32_t dVal)
{
if (dVal < STRIDE_LIMIT_H) {
AscendC::DataCopy(l1Tensor,
gmTensor,
AscendC::Nd2NzParams(1, // ndNum
nTileActual, // nValue
dTileActual, // dValue
0, // srcNdMatrixStride, unused
dVal, // srcDValue
nTileCeil, // dstNzC0Stride
1, // dstNzNStride
0)); // dstNzMatrixStride, unused
} else {
for (uint32_t i = 0; i < nTileActual; ++i) {
AscendC::DataCopy(l1Tensor,
gmTensor,
AscendC::Nd2NzParams(1, // ndNum
1, // nValue
dTileActual, // dValue
0, // srcNdMatrixStride, unused
0, // srcDValue
nTileCeil, // dstNzC0Stride
0, // dstNzNStride
0)); // dstNzMatrixStride, unused
}
}
};
};
#endif // GM_TO_L1_ITERATOR_H

97
csrc/common/include/kernel/gm_to_ub_iterator.h Normal file
View File

@@ -0,0 +1,97 @@
/**
 * Copyright (c) 2025 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 gm_to_ub_iterator.h
* \brief
*/
#ifndef GM_TO_UB_ITERATOR_H
#define GM_TO_UB_ITERATOR_H
#include "iterator.h"
constexpr uint32_t STRIDE_LIMIT_I = 65536;
template <ArchType ArchTag, typename DType> struct gm_to_ub {
__aicore__ inline gm_to_ub(AscendC::LocalTensor<DType> dstTensor, AscendC::GlobalTensor<DType> srcTensor,
uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride)
{
AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride));
};
};
template <ArchType ArchTag, typename DType> struct gm_to_ub_align {
__aicore__ inline gm_to_ub_align(AscendC::LocalTensor<DType> dstTensor, AscendC::GlobalTensor<DType> srcTensor,
uint8_t sid, uint16_t nBurst, uint32_t lenBurst, uint8_t leftPaddingNum,
uint8_t rightPaddingNum, uint32_t srcGap, uint32_t dstGap)
{
AscendC::DataCopyPad(dstTensor, srcTensor, AscendC::DataCopyExtParams(nBurst, lenBurst, srcGap, dstGap, 0),
AscendC::DataCopyPadExtParams<DType>(false, leftPaddingNum, rightPaddingNum, 0));
};
};
template <ArchType ArchTag, typename DType> struct ub_to_ub {
__aicore__ inline ub_to_ub(AscendC::LocalTensor<DType> dstTensor, AscendC::LocalTensor<DType> srcTensor,
uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride)
{
AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride));
};
};
template <ArchType ArchTag, typename DataType, DataFormatT InDataFormat = DataFormatT::ND,
DataFormatT OutDataFormat = DataFormatT::ND>
struct ub_to_gm {
__aicore__ inline ub_to_gm(AscendC::GlobalTensor<DataType> dstTensor, AscendC::LocalTensor<DataType> srcTensor,
uint8_t sid, uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride)
{
AscendC::DataCopy(dstTensor, srcTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride));
};
};
template <ArchType ArchTag, typename DataType> struct ub_to_gm<ArchTag, DataType, DataFormatT::NZ, DataFormatT::NZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
__aicore__ ub_to_gm(AscendC::GlobalTensor<DataType> gmTensor, AscendC::LocalTensor<DataType> l1Tensor,
uint32_t nTileActual, uint32_t nTileCeil, uint32_t nVal, uint32_t dTileActual,
uint32_t dTileCeil, uint32_t dVal)
{
uint64_t dstStride = nTileCeil - nTileActual;
if (dstStride < STRIDE_LIMIT_I) {
AscendC::DataCopy(gmTensor, l1Tensor,
AscendC::DataCopyParams(dTileActual / BLOCK_SIZE, // nBurst
nTileActual, // lenBurst
0, // srcGap
dstStride)); // dstGap
} else {
for (uint64_t i = 0; i < dTileActual / BLOCK_SIZE; i++) {
uint64_t srcOffset = i * nTileActual * BLOCK_SIZE;
uint64_t dstOffset = i * nTileCeil * BLOCK_SIZE;
AscendC::DataCopy(gmTensor[dstOffset], l1Tensor[srcOffset],
AscendC::DataCopyParams(1, // nBurst
nTileActual, // lenBurst
0, // srcGap
0)); // dstGap
}
}
};
};
template <ArchType ArchTag, typename DType> struct ub_to_gm_align {
__aicore__ inline ub_to_gm_align(AscendC::GlobalTensor<DType> dstTensor, AscendC::LocalTensor<DType> srcTensor,
uint8_t sid, uint16_t nBurst, uint32_t lenBurst, uint8_t leftPaddingNum,
uint8_t rightPaddingNum, uint32_t srcGap, uint32_t dstGap)
{
AscendC::DataCopyPad(dstTensor, srcTensor, AscendC::DataCopyExtParams(nBurst, lenBurst, srcGap, dstGap, 0));
};
};
#endif // GM_TO_UB_ITERATOR_H

40
csrc/common/include/kernel/hardware.h Normal file
View File

@@ -0,0 +1,40 @@
/**
 * Copyright (c) 2025 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 hardware.h
* \brief
*/
#ifndef INCLUDE_HARDWARE_H
#define INCLUDE_HARDWARE_H
enum class ArchType { ASCEND_V220, ASCEND_V200, ASCEND_M200 };
template <ArchType ArchTag>
struct HardwareInfo {
static uint32_t const l2BW = 5;
static uint32_t const hbmBW = 1;
static uint32_t const supportMix = 0;
static uint32_t const l1Size = 512 * 1024;
static uint32_t const l0ASize = 64 * 1024;
static uint32_t const l0BSize = 64 * 1024;
static uint32_t const l0CSize = 128 * 1024;
static uint32_t const l2Size = 192 * 1024 * 1024;
static uint32_t const biasSize = 1024;
static uint32_t const fixBufSize = 7 * 1024;
static uint32_t const ubSize = 192 * 1024;
static uint32_t const fractalSize = 512;
static uint32_t const l1l0BlockSize = 32;
static uint32_t const btBlockSize = 64;
static uint32_t const fbBlockSize = 128;
};
#endif

123
csrc/common/include/kernel/iterator.h Normal file
View File

@@ -0,0 +1,123 @@
/**
 * Copyright (c) 2025 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 iterator.h
* \brief
*/
#ifndef INCLUDE_ITERTOR_H
#define INCLUDE_ITERTOR_H
#include "common_func.h"
#include "hardware.h"
#include "kernel_operator.h"
#include "layout.h"
#include "mem.h"
/////////////////////////////////////////////////////
// gm_to_l1
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType, DataFormatT FormatInGM, DataFormatT FormatInL1>
struct gm_to_l1 {
__aicore__ gm_to_l1(AscendC::LocalTensor<DataType> l1Tensor,
AscendC::GlobalTensor<DataType> gmTensor,
uint32_t nTileActual,
uint32_t nTileCeil,
uint32_t nVal,
uint32_t dTileActual,
uint32_t dTileCeil,
uint32_t dVal) {};
};
/////////////////////////////////////////////////////
// l1_to_l0_a
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType, bool IsTransPose, DataFormatT DFmtIn, DataFormatT DFmtOut>
struct l1_to_l0_a {
__aicore__ l1_to_l0_a(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t mTileCeil,
uint32_t kPartCeil,
uint32_t mSrcStride,
uint32_t kSrcStride,
uint32_t mDstStride,
uint32_t kDstStride) {};
};
/////////////////////////////////////////////////////
// l1_to_l0_b
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType, bool IsTransPose, DataFormatT DFmtIn, DataFormatT DFmtOut>
struct l1_to_l0_b {
__aicore__ l1_to_l0_b(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t nTileCeil,
uint32_t kPartCeil,
uint32_t nSrcStride,
uint32_t kSrcStride,
uint32_t nDstStride,
uint32_t kDstStride) {};
};
// l1_to_l0_a
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType, bool IsTransPose, bool IsVectore>
struct l1_to_l0_a_v1 {
__aicore__ l1_to_l0_a_v1(AscendC::LocalTensor<DataType> l0_tensor,
AscendC::LocalTensor<DataType> l1_tensor,
uint32_t m_tile_ceil,
uint32_t k_tile_ceil,
uint32_t k_part,
uint32_t k_part_ceil,
uint32_t k_part_idx) {};
};
/////////////////////////////////////////////////////
// l1_to_l0_b
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType, bool IsTransPose, bool IsVectore>
struct l1_to_l0_b_v1 {
__aicore__ l1_to_l0_b_v1(AscendC::LocalTensor<DataType> l0_tensor,
AscendC::LocalTensor<DataType> l1_tensor,
int32_t n_tile_ceil,
int32_t k_tile_ceil,
int32_t k_part_ceil,
int32_t k_part_idx) {};
};
/////////////////////////////////////////////////////
// l0c_to_gm
/////////////////////////////////////////////////////
template <ArchType ArchTag, DataFormatT OutFormatType, typename OutDataType, typename L0CDataType>
struct l0c_to_gm {
__aicore__ l0c_to_gm(AscendC::GlobalTensor<OutDataType> gmTensor,
AscendC::LocalTensor<L0CDataType> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t mTileCeil,
uint32_t nActual) {};
};
/////////////////////////////////////////////////////
// l0c_to_l1
/////////////////////////////////////////////////////
template <ArchType ArchTag, DataFormatT LayoutOut, typename ElementOut, typename ElementIn>
struct l0c_to_l1 {
__aicore__ l0c_to_l1(AscendC::LocalTensor<ElementOut> l1Tensor,
AscendC::LocalTensor<ElementIn> l0cTensor,
AscendC::LocalTensor<uint64_t> deqTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t mTileCeil,
uint32_t nActual) {};
};
#endif

213
csrc/common/include/kernel/l0c_to_gm_iterator.h Normal file
View File

@@ -0,0 +1,213 @@
/**
 * Copyright (c) 2025 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 l0c_to_gm_iterator.h
* \brief
*/
#ifndef L0C_TO_GM_ITERATOR_H
#define L0C_TO_GM_ITERATOR_H
#ifdef __CCE_KT_TEST__
#define __bf16 bfloat16_t
#endif
#include "iterator.h"
constexpr uint32_t BLOCK_NUM = 16;
constexpr uint32_t BLOCK_SIZE_INT8 = 32;
template <>
struct l0c_to_gm<ArchType::ASCEND_V220, DataFormatT::ND, half, float> {
/**
* @brief Copy data from L0C buffer to global memory, partial specialized for
*
* @param gmTensor the destination tensor on global memory, which is stored in ND format.
* @param l0cTensor the source tensor on L0C buffer, which is stored in FRACTAL_NZ format.
* @param mTileActual the m-direction size of the matrix in L0C buffer.
* @param nTileActual the n-direction size of the matrix in L0C buffer.
* @param srcStride the source stride between the adjacent fractal matrices along n-direction in unit of C0_SIZE.
* @param dstStride the leading dimension of the destination matrix in unit of element.
*/
__aicore__ l0c_to_gm(AscendC::GlobalTensor<half> gmTensor,
AscendC::LocalTensor<float> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t srcStride,
uint32_t dstStride)
{
#ifdef __DAV_C220_CUBE__
auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize
mTileActual, // mSize
srcStride, // srcStride
dstStride, // dstStride
false); // enRelu
intriParams.quantPre = QuantMode_t::F322F16;
AscendC::Fixpipe<half, float, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams);
#else
AscendC::FixpipeParams<float> intriParams(
(nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE,
static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8),
0,
dstStride);
intriParams.nz2ndParams = {true, 1, 0, 0, static_cast<uint16_t>(nTileActual)};
intriParams.quantParams = {QuantMode_t::F322F16};
AscendC::Fixpipe(gmTensor, l0cTensor, intriParams);
#endif
};
};
template <>
struct l0c_to_gm<ArchType::ASCEND_V220, DataFormatT::ND, half, int32_t> {
__aicore__ l0c_to_gm(AscendC::GlobalTensor<half> gmTensor,
AscendC::LocalTensor<int32_t> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t srcStride,
uint32_t dstStride)
{
#ifdef __DAV_C220_CUBE__
auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize
mTileActual, // mSize
srcStride, // srcStride
dstStride, // dstStride
false); // enRelu
intriParams.quantPre = QuantMode_t::VDEQF16;
AscendC::Fixpipe<half, int32_t, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams);
#else
AscendC::FixpipeParams<int32_t> intriParams(
(nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE,
static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8),
0,
dstStride);
intriParams.nz2ndParams = {true, 1, 0, 0, static_cast<uint16_t>(nTileActual)};
intriParams.quantParams = {QuantMode_t::VDEQF16};
AscendC::Fixpipe(gmTensor, l0cTensor, intriParams);
#endif
};
};
template <>
struct l0c_to_gm<ArchType::ASCEND_V220, DataFormatT::ND, __bf16, float> {
__aicore__ l0c_to_gm(AscendC::GlobalTensor<__bf16> gmTensor,
AscendC::LocalTensor<float> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t srcStride,
uint32_t dstStride)
{
#ifdef __DAV_C220_CUBE__
auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize
mTileActual, // mSize
srcStride, // srcStride
dstStride, // dstStride
false); // enRelu
intriParams.quantPre = QuantMode_t::F322BF16;
AscendC::Fixpipe<__bf16, float, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams);
#else
AscendC::FixpipeParams<float> intriParams(
(nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE,
static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8),
0,
dstStride);
intriParams.nz2ndParams = {true, 1, 0, 0, static_cast<uint16_t>(nTileActual)};
intriParams.quantParams = {QuantMode_t::F322BF16};
AscendC::Fixpipe(gmTensor, l0cTensor, intriParams);
#endif
};
};
// Partial specialization ND, float
template <>
struct l0c_to_gm<ArchType::ASCEND_V220, DataFormatT::ND, float, float> {
__aicore__ l0c_to_gm(AscendC::GlobalTensor<float> gmTensor,
AscendC::LocalTensor<float> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t srcStride,
uint32_t dstStride)
{
#ifdef __DAV_C220_CUBE__
auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize
mTileActual, // mSize
srcStride, // srcStride
dstStride, // dstStride
false); // enRelu
intriParams.quantPre = QuantMode_t::NoQuant;
AscendC::Fixpipe<float, float, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams);
#else
AscendC::FixpipeParams<float> intriParams(
(nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE,
static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8),
0,
dstStride);
intriParams.nz2ndParams = {true, 1, 0, 0, static_cast<uint16_t>(nTileActual)};
intriParams.quantParams = {QuantMode_t::NoQuant};
AscendC::Fixpipe(gmTensor, l0cTensor, intriParams);
#endif
};
};
template <>
struct l0c_to_gm<ArchType::ASCEND_V220, DataFormatT::NZ, half, float> {
__aicore__ l0c_to_gm(AscendC::GlobalTensor<half> gmTensor,
AscendC::LocalTensor<float> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t srcStride,
uint32_t dstStride)
{
#ifdef __DAV_C220_CUBE__
auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize
mTileActual, // mSize
srcStride, // srcStride
dstStride, // dstStride
false); // enRelu
intriParams.quantPre = QuantMode_t::F322F16;
AscendC::Fixpipe<half, float, AscendC::CFG_NZ>(gmTensor, l0cTensor, intriParams);
#else
AscendC::FixpipeParams<float> intriParams(
(nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE,
static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE_INT8),
0,
dstStride - (nTileActual * sizeof(half) / sizeof(float)));
intriParams.quantParams = {QuantMode_t::F322F16};
AscendC::Fixpipe(gmTensor, l0cTensor, intriParams);
#endif
};
};
template <>
struct l0c_to_gm<ArchType::ASCEND_V220, DataFormatT::ND, int32_t, int32_t> {
__aicore__ l0c_to_gm(AscendC::GlobalTensor<int32_t> gmTensor,
AscendC::LocalTensor<int32_t> l0cTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t srcStride,
uint32_t dstStride){
#ifdef __DAV_C220_CUBE__
auto intriParams = AscendC::FixpipeParamsV220(nTileActual, // nSize
mTileActual, // mSize
srcStride, // srcStride
dstStride, // dstStride
false); // enRelu
intriParams.quantPre = QuantMode_t::NoQuant;
AscendC::Fixpipe<int32_t, int32_t, AscendC::CFG_ROW_MAJOR>(gmTensor, l0cTensor, intriParams);
#endif
};
};
#endif // L0C_TO_GM_ITERATOR_H

51
csrc/common/include/kernel/l0c_to_l1_iterator.h Normal file
View File

@@ -0,0 +1,51 @@
/**
 * Copyright (c) 2025 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 l0c_to_l1_iterator.h
* \brief
*/
#ifndef L0C_TO_L1_ITERATOR_H
#define L0C_TO_L1_ITERATOR_H
#include "iterator.h"
/////////////////////////////////////////////////////
// l0c_to_l1
/////////////////////////////////////////////////////
// Partial specialization ZN, half, int32_t
template <ArchType ArchTag>
struct l0c_to_l1<ArchTag, DataFormatT::ZN, half, int32_t> {
using ElementOut = half;
using ElementIn = int32_t;
__aicore__ l0c_to_l1(AscendC::LocalTensor<ElementOut> l1Tensor,
AscendC::LocalTensor<ElementIn> l0cTensor,
AscendC::LocalTensor<uint64_t> deqTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t mTileCeil,
uint32_t nActual)
{
constexpr uint32_t BLOCK_NUM = 16;
constexpr uint32_t BLOCK_SIZE = 32;
AscendC::FixpipeParams<ElementIn> intriParams(
(nTileActual + BLOCK_NUM - 1) / AscendC::BLOCK_CUBE,
static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE),
0,
mTileCeil - static_cast<uint16_t>(mTileActual * BLOCK_NUM * sizeof(float) / BLOCK_SIZE) *
sizeof(ElementOut) / sizeof(ElementIn));
intriParams.nz2ndParams = {false, 1, 0, 0, static_cast<uint16_t>(nTileActual)};
intriParams.quantParams = {QuantMode_t::VDEQF16};
AscendC::Fixpipe(l1Tensor, l0cTensor, deqTensor, intriParams);
};
};
#endif // L0C_TO_L1_ITERATOR_H

72
csrc/common/include/kernel/l0c_to_ub_iterator.h Normal file
View File

@@ -0,0 +1,72 @@
/**
 * Copyright (c) 2025 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 l0c_to_ub_iterator.h
* \brief
*/
#ifndef L0C_TO_UB_ITERATOR_H
#define L0C_TO_UB_ITERATOR_H
#include "iterator.h"
/////////////////////////////////////////////////////
// l0c_to_ub
/////////////////////////////////////////////////////
// Partial specialization ZN, half, int32_t
template <ArchType ArchTag, typename ElementIn, typename ElementOut, bool MatrixMode = true> struct l0c_to_ub {
__aicore__ l0c_to_ub(AscendC::LocalTensor<ElementOut> ubTensor, AscendC::LocalTensor<ElementIn> l0cTensor,
uint16_t nBurst, uint16_t lenBurst, uint16_t srcStride, uint16_t dstStride)
{
constexpr auto mode =
MatrixMode ? AscendC::BlockMode::BLOCK_MODE_MATRIX : AscendC::BlockMode::BLOCK_MODE_VECTOR;
AscendC::DataCopy(ubTensor, l0cTensor,
AscendC::DataCopyParams(nBurst, // count
lenBurst, // len
srcStride, // srcStrideIn
dstStride), // dstStrideIn
AscendC::DataCopyEnhancedParams(mode, // blockModeIn
AscendC::DeqScale::DEQ_NONE, // deqScaleIn
0, // deqValueIn
0, // sidStoreModeIn
false, // isReluIn
pad_t::PAD_NONE, // padModeIn
0) // padValueIn
);
};
};
template <ArchType ArchTag>
struct l0c_to_ub<ArchTag, int32_t, half> {
__aicore__ l0c_to_ub(AscendC::LocalTensor<half> ubTensor,
AscendC::LocalTensor<int32_t> l0cTensor,
uint16_t nBurst,
uint16_t lenBurst,
uint16_t srcStride,
uint16_t dstStride)
{
AscendC::DataCopy(ubTensor, l0cTensor,
AscendC::DataCopyParams(nBurst, // count
lenBurst, // len
srcStride, // srcStrideIn
dstStride), // dstStrideIn
AscendC::DataCopyEnhancedParams(AscendC::BlockMode::BLOCK_MODE_MATRIX, // blockModeIn
AscendC::DeqScale::VDEQ16, // deqScaleIn
0, // deqValueIn
0, // sidStoreModeIn
false, // isReluIn
pad_t::PAD_NONE, // padModeIn
0) // padValueIn
);
};
};
#endif // L0C_TO_UB_ITERATOR_H

39
csrc/common/include/kernel/l1_to_bt_iterator.h Normal file
View File

@@ -0,0 +1,39 @@
/**
 * Copyright (c) 2025 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 l1_to_bt_iterator.h
* \brief
*/
#ifndef L1_TO_BT_ITERATOR_H
#define L1_TO_BT_ITERATOR_H
#include "iterator.h"
/////////////////////////////////////////////////////
// l1_to_bt
/////////////////////////////////////////////////////
// Partial specialization for V220
template <ArchType ArchTag, typename DataType>
struct l1_to_bt {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::btBlockSize / sizeof(DataType);
__aicore__ l1_to_bt(AscendC::LocalTensor<DataType> biasTableTensor,
AscendC::LocalTensor<DataType> biasL1Tensor,
uint32_t ntileActual)
{
AscendC::DataCopy(
biasTableTensor, biasL1Tensor, {1, static_cast<uint16_t>(CeilDiv<BLOCK_SIZE>(ntileActual)), 0, 0});
};
};
#endif // L1_TO_BT_ITERATOR_H

43
csrc/common/include/kernel/l1_to_fb_iterator.h Normal file
View File

@@ -0,0 +1,43 @@
/**
 * Copyright (c) 2025 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 l1_to_fb_iterator.h
* \brief
*/
#ifndef L1_TO_FB_ITERATOR_H
#define L1_TO_FB_ITERATOR_H
#include "iterator.h"
/////////////////////////////////////////////////////
// l1_to_fb
/////////////////////////////////////////////////////
// Partial specialization for V220
template <ArchType ArchTag, typename DataType>
struct l1_to_fb {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::fbBlockSize / sizeof(DataType);
__aicore__
l1_to_fb(AscendC::LocalTensor<DataType> fbTensor, AscendC::LocalTensor<DataType> l1Tensor, uint32_t ntileActual)
{
copy_cbuf_to_fbuf((__fbuf__ DataType *)fbTensor.GetPhyAddr(),
(__cbuf__ DataType *)l1Tensor.GetPhyAddr(),
1,
CeilDiv<BLOCK_SIZE>(ntileActual),
0,
0);
};
};
#endif // L1_TO_FB_ITERATOR_H

259
csrc/common/include/kernel/l1_to_l0_iterator.h Normal file
View File

@@ -0,0 +1,259 @@
/**
 * Copyright (c) 2025 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 l1_to_l0_iterator.h
* \brief
*/
#ifndef L1_TO_L0_ITERATOR_H
#define L1_TO_L0_ITERATOR_H
#include "iterator.h"
/////////////////////////////////////////////////////
// l1_to_l0_a
/////////////////////////////////////////////////////
// Partial specialization for vector
template <ArchType ArchTag, typename DataType, bool IsTransPose>
struct l1_to_l0_a<ArchTag, DataType, IsTransPose, DataFormatT::VECTOR, DataFormatT::VECTOR> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
__aicore__ l1_to_l0_a(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t mTileCeil,
uint32_t kPartCeil,
uint32_t mSrcStride,
uint32_t kSrcStride,
uint32_t mDstStride,
uint32_t kDstStride)
{
AscendC::LoadData(l0Tensor,
l1Tensor,
AscendC::LoadData2dParams(0, // baseIdx
kPartCeil, // repeat
kSrcStride, // srcStride
0, // sid
kDstStride, // dstStride
IsTransPose, // transpose
0)); // addrCalMode
};
};
// Partial specialization for no transpose, not vector
template <ArchType ArchTag, typename DataType>
struct l1_to_l0_a<ArchTag, DataType, false, DataFormatT::ZN, DataFormatT::ZZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize;
__aicore__ l1_to_l0_a(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t mTileCeil,
uint32_t kPartCeil,
uint32_t mSrcStride,
uint32_t kSrcStride,
uint32_t mDstStride,
uint32_t kDstStride)
{
for (uint32_t i = 0; i < mTileCeil / BLOCK_NUM_PER_FRACTAL; i++) {
AscendC::LoadData(l0Tensor[i * mDstStride * FRACTAL_SIZE],
l1Tensor[i * mSrcStride * FRACTAL_SIZE],
AscendC::LoadData2dParams(0, // baseIdx
static_cast<uint16_t>(kPartCeil / BLOCK_SIZE), // repeat
kSrcStride, // srcStride
0, // sid
kDstStride - 1, // dstStride
false, // transpose
0)); // addrCalMode
}
};
};
// Partial specialization for transpose, not vector
template <ArchType ArchTag, typename DataType>
struct l1_to_l0_a<ArchTag, DataType, true, DataFormatT::ZN, DataFormatT::ZZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize;
__aicore__ l1_to_l0_a(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t mTileCeil,
uint32_t kPartCeil,
uint32_t mSrcStride,
uint32_t kSrcStride,
uint32_t mDstStride,
uint32_t kDstStride)
{
for (uint32_t i = 0; i < mTileCeil / BLOCK_SIZE; i++) {
AscendC::LoadData(l0Tensor[i * mDstStride * FRACTAL_SIZE],
l1Tensor[i * mSrcStride * FRACTAL_SIZE],
AscendC::LoadData2dParams(0,
static_cast<uint16_t>(kPartCeil / BLOCK_NUM_PER_FRACTAL),
kSrcStride,
0,
kDstStride - 1,
true,
0));
}
};
};
template <ArchType ArchTag, typename DataType>
struct l1_to_l0_a<ArchTag, DataType, false, DataFormatT::NZ, DataFormatT::ZZ> {
using HardwareParams = HardwareInfo<ArchTag>;
// 16 * 32
static constexpr uint32_t ROW_BLOCK_SIZE = 16;
static constexpr uint32_t COL_BLOCK_SIZE = 32 / sizeof(DataType);
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize;
__aicore__ l1_to_l0_a(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t mTileCeil,
uint32_t kPartCeil,
uint32_t mSrcStride,
uint32_t kSrcStride,
uint32_t mDstStride,
uint32_t kDstStride)
{
for (uint32_t i = 0; i < mTileCeil / ROW_BLOCK_SIZE; i++) {
AscendC::LoadData(l0Tensor[i * ROW_BLOCK_SIZE * kPartCeil],
l1Tensor[i * FRACTAL_SIZE],
AscendC::LoadData2dParams(0,
static_cast<uint16_t>(kPartCeil / COL_BLOCK_SIZE),
mTileCeil / ROW_BLOCK_SIZE,
0,
0,
false,
0));
}
};
};
/////////////////////////////////////////////////////
// l1_to_l0_b
/////////////////////////////////////////////////////
// Partial specialization for vector
template <ArchType ArchTag, typename DataType, bool IsTransPose>
struct l1_to_l0_b<ArchTag, DataType, IsTransPose, DataFormatT::VECTOR, DataFormatT::VECTOR> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
__aicore__ l1_to_l0_b(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t nTileCeil,
uint32_t kPartCeil,
uint32_t nSrcStride,
uint32_t kSrcStride,
uint32_t nDstStride,
uint32_t kDstStride)
{
AscendC::LoadData(
l0Tensor, l1Tensor, AscendC::LoadData2dParams(0, kPartCeil, kSrcStride, 0, kDstStride, IsTransPose, 0));
};
};
template <ArchType ArchTag>
struct l1_to_l0_b<ArchTag, int8_t, true, DataFormatT::NZ, DataFormatT::ZN> {
using HardwareParams = HardwareInfo<ArchTag>;
using DataType = int8_t;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
__aicore__ l1_to_l0_b(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t nTileCeil,
uint32_t kPartCeil,
uint32_t nSrcStride,
uint32_t kSrcStride,
uint32_t nDstStride,
uint32_t kDstStride)
{
for (uint32_t i = 0; i < nTileCeil / BLOCK_SIZE; i++) {
AscendC::LoadDataWithTranspose(l0Tensor[i * kPartCeil * BLOCK_SIZE],
l1Tensor[i * BLOCK_SIZE * BLOCK_SIZE],
AscendC::LoadData2dTransposeParams(0, // startIndexIn
kPartCeil / BLOCK_SIZE, // repeatTimesIn
nTileCeil / BLOCK_SIZE, // srcStrideIn
1, // dstGapIn
0, // dstfracGapIn
0) // addrModeIn
);
}
};
};
// Partial specialization for no transpose, not vector
template <ArchType ArchTag, typename DataType>
struct l1_to_l0_b<ArchTag, DataType, false, DataFormatT::ZN, DataFormatT::NZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize;
__aicore__ l1_to_l0_b(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t nTileCeil,
uint32_t kPartCeil,
uint32_t nSrcStride,
uint32_t kSrcStride,
uint32_t nDstStride,
uint32_t kDstStride)
{
for (uint32_t i = 0; i < kPartCeil / BLOCK_NUM_PER_FRACTAL; i++) {
AscendC::LoadData(l0Tensor[i * kDstStride * FRACTAL_SIZE],
l1Tensor[i * kSrcStride * FRACTAL_SIZE],
AscendC::LoadData2dParams(0, // baseIdx
static_cast<uint16_t>(nTileCeil / BLOCK_SIZE), // repeat
nSrcStride, // srcStride
0, // sid
nDstStride - 1, // dstStride
true, // transpose
0)); // addrCalMode
}
};
};
// Partial specialization for transpose, not vector
template <ArchType ArchTag, typename DataType>
struct l1_to_l0_b<ArchTag, DataType, true, DataFormatT::ZN, DataFormatT::NZ> {
using HardwareParams = HardwareInfo<ArchTag>;
static constexpr uint32_t BLOCK_SIZE = HardwareParams::l1l0BlockSize / sizeof(DataType);
static constexpr uint32_t FRACTAL_SIZE = HardwareParams::fractalSize / sizeof(DataType);
static constexpr uint32_t BLOCK_NUM_PER_FRACTAL = HardwareParams::fractalSize / HardwareParams::l1l0BlockSize;
__aicore__ l1_to_l0_b(AscendC::LocalTensor<DataType> l0Tensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint32_t nTileCeil,
uint32_t kPartCeil,
uint32_t nSrcStride,
uint32_t kSrcStride,
uint32_t nDstStride,
uint32_t kDstStride)
{
AscendC::LoadData(
l0Tensor,
l1Tensor,
AscendC::LoadData2dParams(0, // baseIdx
static_cast<uint16_t>(kPartCeil * nTileCeil / FRACTAL_SIZE), // repeat
1, // srcStride
0, // sid
0, // dstStride
false, // transpose
0)); // addr_cal_mode_t
};
};
#endif // L1_TO_L0_ITERATOR_H

52
csrc/common/include/kernel/l1_to_ub_iterator.h Normal file
View File

@@ -0,0 +1,52 @@
/**
 * Copyright (c) 2025 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 l1_to_ub_iterator.h
* \brief
*/
#ifndef L1_TO_UB_ITERATOR_H
#define L1_TO_UB_ITERATOR_H
#include "iterator.h"
/////////////////////////////////////////////////////
// l1_to_ub
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType>
struct l1_to_ub {
__aicore__ l1_to_ub(AscendC::LocalTensor<DataType> ubTensor,
AscendC::LocalTensor<DataType> l1Tensor,
uint16_t nBurst,
uint16_t lenBurst,
uint16_t srcStride,
uint16_t dstStride)
{
AscendC::DataCopy(ubTensor, l1Tensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride));
};
};
/////////////////////////////////////////////////////
// ub_to_l1
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DataType>
struct ub_to_l1 {
__aicore__ ub_to_l1(AscendC::LocalTensor<DataType> l1Tensor,
AscendC::LocalTensor<DataType> ubTensor,
uint16_t nBurst,
uint16_t lenBurst,
uint16_t srcStride,
uint16_t dstStride)
{
AscendC::DataCopy(l1Tensor, ubTensor, AscendC::DataCopyParams(nBurst, lenBurst, srcStride, dstStride));
};
};
#endif // L1_TO_UB_ITERATOR_H

28
csrc/common/include/kernel/layout.h Normal file
View File

@@ -0,0 +1,28 @@
/**
 * Copyright (c) 2025 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 layout.h
* \brief
*/
#ifndef INCLUDE_LAYOUT_H
#define INCLUDE_LAYOUT_H
enum class DataFormatT {
ND = 0,
NZ,
ZN,
ZZ,
NN,
VECTOR
};
#endif // INCLUDE_LAYOUT_H

79
csrc/common/include/kernel/mem.h Normal file
View File

@@ -0,0 +1,79 @@
/**
 * Copyright (c) 2025 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 mem.h
* \brief
*/
#ifndef INCLUDE_MEM_H
#define INCLUDE_MEM_H
#include "hardware.h"
#include "kernel_event.h"
#include "kernel_tensor.h"
enum class BufferType { ASCEND_UB, ASCEND_CB, ASCEND_L0A, ASCEND_L0B, ASCEND_L0C, ASCEND_MAX };
template <BufferType BufferType_>
__aicore__ constexpr AscendC::TPosition GetPosition()
{
if constexpr (BufferType_ == BufferType::ASCEND_UB) {
return AscendC::TPosition::VECIN;
} else if constexpr (BufferType_ == BufferType::ASCEND_CB) {
return AscendC::TPosition::A1;
} else if constexpr (BufferType_ == BufferType::ASCEND_L0A) {
return AscendC::TPosition::A2;
} else if constexpr (BufferType_ == BufferType::ASCEND_L0B) {
return AscendC::TPosition::B2;
} else if constexpr (BufferType_ == BufferType::ASCEND_L0C) {
return AscendC::TPosition::CO1;
}
return AscendC::TPosition::GM;
}
template <ArchType ArchTag>
struct AsdopsBuffer {
public:
__aicore__ AsdopsBuffer()
{
constexpr uint32_t bufferSize[(uint32_t)BufferType::ASCEND_MAX] = {HardwareInfo<ArchTag>::ubSize,
HardwareInfo<ArchTag>::l1Size,
HardwareInfo<ArchTag>::l0ASize,
HardwareInfo<ArchTag>::l0BSize,
HardwareInfo<ArchTag>::l0CSize};
#ifdef __DAV_C220_VEC__
tensor[(uint32_t)BufferType::ASCEND_UB] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::VECIN, 0, bufferSize[(uint32_t)BufferType::ASCEND_UB]);
#elif __DAV_C220_CUBE__
tensor[(uint32_t)BufferType::ASCEND_CB] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::A1, 0, bufferSize[(uint32_t)BufferType::ASCEND_CB]);
tensor[(uint32_t)BufferType::ASCEND_L0A] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::A2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]);
tensor[(uint32_t)BufferType::ASCEND_L0B] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::B2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]);
tensor[(uint32_t)BufferType::ASCEND_L0C] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::CO1, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]);
#else
#ifndef __clang__
tensor[(uint32_t)BufferType::ASCEND_UB] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::VECIN, 0, bufferSize[(uint32_t)BufferType::ASCEND_UB]);
tensor[(uint32_t)BufferType::ASCEND_CB] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::A1, 0, bufferSize[(uint32_t)BufferType::ASCEND_CB]);
tensor[(uint32_t)BufferType::ASCEND_L0A] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::A2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0A]);
tensor[(uint32_t)BufferType::ASCEND_L0B] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::B2, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0B]);
tensor[(uint32_t)BufferType::ASCEND_L0C] = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::CO1, 0, bufferSize[(uint32_t)BufferType::ASCEND_L0C]);
#endif
#endif
};
template <BufferType BufferType_, typename DstDataType = half>
__aicore__ AscendC::LocalTensor<DstDataType> GetBuffer(const uint32_t offset) const
{
return tensor[(uint32_t)BufferType_][offset].template ReinterpretCast<DstDataType>();
}
public:
AscendC::LocalTensor<uint8_t> tensor[(uint32_t)BufferType::ASCEND_MAX];
};
#endif

79
csrc/common/include/kernel/mma.h Normal file
View File

@@ -0,0 +1,79 @@
/**
 * Copyright (c) 2025 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 mma.h
* \brief
*/
#ifndef INCLUDE_MMA_H
#define INCLUDE_MMA_H
#include "hardware.h"
#include "kernel_tensor.h"
template <ArchType ArchTag, typename ElementA, typename ElementB, typename AccDTypeC, bool IsTransposeA>
struct mmad {
__aicore__ mmad(AscendC::LocalTensor<AccDTypeC> l0cTensor,
AscendC::LocalTensor<ElementA> l0aTensor,
AscendC::LocalTensor<ElementB> l0bTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t kPartActual,
bool initC) {};
__aicore__ mmad(AscendC::LocalTensor<AccDTypeC> l0cTensor,
AscendC::LocalTensor<ElementA> l0aTensor,
AscendC::LocalTensor<ElementB> l0bTensor,
uint64_t biasBt,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t kPartActual,
bool initC) {};
};
// Partial specialization for V220, int8_t, not_vector_A, not TransposeA
template <ArchType ArchTag, typename AccDTypeC, typename ElementA, typename ElementB>
struct mmad<ArchTag, ElementA, ElementB, AccDTypeC, false> {
__aicore__ mmad(AscendC::LocalTensor<AccDTypeC> l0cTensor,
AscendC::LocalTensor<ElementA> l0aTensor,
AscendC::LocalTensor<ElementB> l0bTensor,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t kPartActual,
bool initC)
{
AscendC::Mmad(l0cTensor,
l0aTensor,
l0bTensor,
AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC));
};
__aicore__ mmad(AscendC::LocalTensor<AccDTypeC> l0cTensor,
AscendC::LocalTensor<ElementA> l0aTensor,
AscendC::LocalTensor<ElementB> l0bTensor,
uint64_t biasBt,
uint32_t mTileActual,
uint32_t nTileActual,
uint32_t kPartActual,
bool initC)
{
AscendC::LocalTensor<ElementA> biasTensor;
biasTensor.InitBuffer(biasBt, mTileActual);
biasTensor.address_.logicPos = static_cast<uint8_t>(AscendC::TPosition::C2);
AscendC::Mmad(l0cTensor,
l0aTensor,
l0bTensor,
biasTensor,
AscendC::MmadParams(mTileActual, nTileActual, kPartActual, 0, false, initC));
};
};
#endif

483
csrc/common/include/kernel/pse.h Normal file
View File

@@ -0,0 +1,483 @@
/**
 * Copyright (c) 2025 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 pse.h
* \brief
*/
#ifndef FLASH_ATTENTION_SCORE_PSE_H
#define FLASH_ATTENTION_SCORE_PSE_H
#include "kernel_operator.h"
#include "util.h"
constexpr static int64_t pseS1S2 = 0;
constexpr static int64_t pse1S2 = 1;
constexpr static int64_t pseSlopeBn = 2;
constexpr static int64_t pseSlopeN = 3;
constexpr static uint8_t pseEncodeALibiS2Full = 0x11;
enum class PseTypeEnum {
PSE_OUTER_MUL_ADD_TYPE = 0, // default
PSE_OUTER_ADD_MUL_TYPE,
PSE_INNER_MUL_ADD_TYPE,
PSE_INNER_MUL_ADD_SQRT_TYPE,
PSE_INVALID_TYPE
};
struct PseInfo {
int64_t blockCount;
int64_t bSSOffset; // boidx * s1 * s2
int64_t boIdx;
int64_t gSize;
int64_t goIdx;
int64_t loopIdx;
int64_t n2G;
int64_t n2oIdx;
int64_t pseBSize;
int64_t pseS1Size; // for alibi
int64_t pseS2ComputeSize; // for alibi, do not need assignment
int64_t pseS2Size; // for alibi
uint32_t pseShapeType;
int64_t readS2Size; // for alibi, do not need assignment
int64_t s1BaseSize;
int64_t s1Size;
int64_t s1oIdx;
int64_t s2AlignedSize;
int64_t s2BaseNratioSize;
int64_t s2LoopCount;
int64_t s2RealSize;
int64_t s2Size;
int64_t s2SizeAcc; // accumulated sum of s2 size
int64_t s2StartIdx;
int64_t vec1S1BaseSize;
int64_t vec1S1RealSize;
uint32_t pseEncodeType; // for distinguish alibi
uint32_t pseType; // 0: outer, mul-add 1:outer, add-mul 2:inner, mul-add 3:inner, mul-add-sqrt
int64_t pseAlibiBaseS1;
int64_t pseAlibiBaseS2;
int64_t qStartIdx;
int64_t kvStartIdx;
int64_t vecCoreOffset = 0;
bool needCast;
bool align8 = false;
bool pseEndogenous = false;
};
template <typename INPUT_T, bool hasPse>
__aicore__ inline void DataCopyInCommon(LocalTensor<INPUT_T> &dstTensor, GlobalTensor<INPUT_T> &srcTensor, int64_t offset,
int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int32_t dtypeSize,
int32_t alignedS2Size)
{
if constexpr (hasPse == true) {
uint32_t shapeArray[] = {static_cast<uint32_t>(s1Size), static_cast<uint32_t>(alignedS2Size)};
dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND));
dstTensor.SetSize(s1Size * alignedS2Size);
DataCopyParams dataCopyParams;
dataCopyParams.blockCount = s1Size;
dataCopyParams.blockLen = CeilDiv(s2Size * dtypeSize, blockBytes); // 单位32B
dataCopyParams.dstStride = alignedS2Size * dtypeSize / blockBytes - dataCopyParams.blockLen; // gap
if (actualS2Len * dtypeSize % blockBytes == 0) {
dataCopyParams.srcStride =
(actualS2Len * dtypeSize - dataCopyParams.blockLen * blockBytes) / blockBytes; // srcGap
DataCopy(dstTensor, srcTensor[offset], dataCopyParams);
} else {
dataCopyParams.blockLen = s2Size * dtypeSize; // 单位Byte
dataCopyParams.srcStride = (actualS2Len * dtypeSize - dataCopyParams.blockLen);
dataCopyParams.dstStride = (alignedS2Size - s2Size) * dtypeSize / blockBytes;
DataCopyPadParams dataCopyPadParams;
dataCopyPadParams.isPad = false;
DataCopyPad(dstTensor, srcTensor[offset], dataCopyParams, dataCopyPadParams);
}
}
}
template <typename INPUT_T, bool hasPse>
__aicore__ inline void DataCopyIn(LocalTensor<INPUT_T> &dstTensor, GlobalTensor<INPUT_T> &srcTensor, int64_t offset,
int64_t s1Size, int64_t s2Size, int64_t actualS2Len, int64_t alignedSize = 16)
{
if constexpr (hasPse == true) {
int32_t dtypeSize = sizeof(INPUT_T);
int32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize;
DataCopyInCommon<INPUT_T, hasPse>(dstTensor, srcTensor, offset, s1Size, s2Size,
actualS2Len, dtypeSize, alignedS2Size);
}
}
template <typename INPUT_T, bool hasPse>
__aicore__ inline void DataCopyInAlign8(LocalTensor<INPUT_T> &dstTensor, GlobalTensor<INPUT_T> &srcTensor, int64_t offset,
int64_t s1Size, int64_t s2Size, int64_t actualS2Len)
{
if constexpr (hasPse == true) {
int32_t dtypeSize = sizeof(INPUT_T);
if (dtypeSize == 0){
return;
}
int32_t alignedS2Size = CeilDiv(s2Size, 32 / dtypeSize) * (32 / dtypeSize);
DataCopyInCommon<INPUT_T, hasPse>(dstTensor, srcTensor, offset, s1Size, s2Size,
actualS2Len, dtypeSize, alignedS2Size);
}
}
/*
dst = BroadcastAdd(src0, src1)
src0 shape: (s1, s2)
src1 shape: (1, s2)
dst shape: (s1, s2)
*/
template <typename T, bool hasPse>
__aicore__ inline void BroadcastAdd(const LocalTensor<T> &src0Tensor, const LocalTensor<T> &src1Tensor,
int64_t src0Offset, int32_t src1Size, int32_t repeatTimes)
{
if constexpr (hasPse == true) {
/* Total data number of single step should be smaller than 256bytes.
* If larger, we need to do add multiple times. */
int32_t innerLoop = src1Size / repeatMaxSize; // s2轴整块计算次数
int32_t innerRemain = src1Size % repeatMaxSize; // s2轴尾块计算量
BinaryRepeatParams binaryRepeatParams;
binaryRepeatParams.src0BlkStride = 1;
binaryRepeatParams.src0RepStride = src1Size / blockSize;
binaryRepeatParams.src1BlkStride = 1;
binaryRepeatParams.src1RepStride = 0;
binaryRepeatParams.dstRepStride = binaryRepeatParams.src0RepStride;
binaryRepeatParams.blockNumber = binaryRepeatParams.src0RepStride;
for (int32_t j = 0; j < innerLoop; j++) {
auto innerOffset = j * repeatMaxSize;
auto ubOffset = src0Offset + innerOffset;
Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], repeatMaxSize, repeatTimes,
binaryRepeatParams);
}
if (innerRemain > 0) {
auto innerOffset = innerLoop * repeatMaxSize;
auto ubOffset = src0Offset + innerOffset;
Add(src0Tensor[ubOffset], src0Tensor[ubOffset], src1Tensor[innerOffset], innerRemain, repeatTimes,
binaryRepeatParams);
}
}
}
template <typename T, bool hasPse>
__aicore__ inline void PseBroadcastAdd(int32_t s1Size, int32_t s2Size, int32_t computeSize, const LocalTensor<T> &pseUb,
const LocalTensor<T> &dstTensor, uint32_t pseShapeType)
{
if constexpr (hasPse == true) {
if (pseShapeType == pseS1S2 || pseShapeType == pseSlopeBn || pseShapeType == pseSlopeN) {
Add(dstTensor, dstTensor, pseUb, computeSize);
} else {
/* Total repeated times should be <= repeatMaxTimes. If larger,
* we need to do multiple inner loops. */
int32_t s1OuterLoop = s1Size / repeatMaxTimes;
int32_t s1OuterRemain = s1Size % repeatMaxTimes;
for (int32_t s1OuterIdx = 0; s1OuterIdx < s1OuterLoop; s1OuterIdx++) {
int32_t s1OuterOffset = s1OuterIdx * repeatMaxTimes * s2Size;
BroadcastAdd<T, hasPse>(dstTensor, pseUb, s1OuterOffset, s2Size, repeatMaxTimes);
}
if (s1OuterRemain > 0) {
int32_t s1OuterOffset = s1OuterLoop * repeatMaxTimes * s2Size;
BroadcastAdd<T, hasPse>(dstTensor, pseUb, s1OuterOffset, s2Size, s1OuterRemain);
}
}
}
}
template <bool hasPse> __aicore__ inline int64_t PseComputeOffset(PseInfo &pseInfo)
{
if constexpr (hasPse == true) {
int64_t bOffset = 0;
int64_t n2Offset = 0;
int64_t s1Offset = 0;
int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize;
int64_t gOffset = 0;
if (pseInfo.pseShapeType == pseS1S2) {
// b, n2, g, s1, s2
bOffset = pseInfo.bSSOffset * pseInfo.n2G;
n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s1Size * pseInfo.s2Size;
gOffset = pseInfo.goIdx * pseInfo.s1Size * pseInfo.s2Size;
s1Offset = (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset +
pseInfo.loopIdx * pseInfo.vec1S1BaseSize) * pseInfo.s2Size;
} else if (pseInfo.pseShapeType == pse1S2) {
// b, n2, g, 1, s2
bOffset = pseInfo.s2SizeAcc * pseInfo.n2G;
n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.s2Size;
gOffset = pseInfo.goIdx * pseInfo.s2Size;
}
if (pseInfo.pseBSize == 1) {
bOffset = 0;
}
return bOffset + n2Offset + gOffset + s1Offset + s2Offset;
} else {
return 0;
}
}
template <LayOutTypeEnum layOutType, bool hasPse> __aicore__ inline int64_t PseAlibiComputeOffset(PseInfo &pseInfo)
{
if constexpr (hasPse == true) {
int64_t bOffset = (pseInfo.boIdx % pseInfo.pseBSize) * pseInfo.n2G * pseInfo.pseS2Size * pseInfo.pseS1Size;
int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize * pseInfo.pseS2Size * pseInfo.pseS1Size;
int64_t gOffset = pseInfo.goIdx * pseInfo.pseS2Size * pseInfo.pseS1Size;
int64_t row = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset +
pseInfo.loopIdx * pseInfo.vec1S1BaseSize;
int64_t column = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize;
int64_t m = 0;
int64_t k = 0;
if constexpr (layOutType != LayOutTypeEnum::LAYOUT_TND) {
int64_t threshold = pseInfo.s1Size - pseInfo.pseS1Size;
if (row >= threshold) {
m = row - threshold;
k = column;
} else {
m = row % pseInfo.pseS1Size;
k = pseInfo.pseS2Size - (row - column) - (pseInfo.pseS1Size - m);
}
} else {
int64_t threshold = pseInfo.pseS2Size - pseInfo.pseS1Size;
int64_t posVal = row - column - threshold;
if (threshold >= 0) {
if (posVal >= 0) {
m = posVal;
k = 0;
} else {
m = 0;
k = -posVal;
}
} else {
m = posVal;
k = 0;
}
}
int64_t s1Offset = m * pseInfo.pseS2Size;
int64_t s2Offset = k;
pseInfo.readS2Size = Min(pseInfo.s2AlignedSize, pseInfo.pseS2Size - k);
pseInfo.pseS2ComputeSize = Align(pseInfo.readS2Size);
return bOffset + n2Offset + gOffset + s1Offset + s2Offset;
} else {
return 0;
}
}
template <bool hasPse> __aicore__ inline bool NeedPseAlibiCompute(PseInfo &pseInfo)
{
if constexpr (hasPse == true) {
// Alibi编码只计算下三角
if (pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset +
(pseInfo.loopIdx + 1) * pseInfo.vec1S1BaseSize <=
pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize) {
return false;
}
return true;
} else {
return false;
}
}
template <typename INPUT_T, typename T, LayOutTypeEnum layOutType, bool hasPse>
__aicore__ inline void PseAlibiCopyIn(LocalTensor<T> &dstTensor, LocalTensor<INPUT_T> &tmpTensor,
GlobalTensor<INPUT_T> &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16)
{
if constexpr (hasPse == true) {
if (!NeedPseAlibiCompute<hasPse>(pseInfo)) {
return;
}
int64_t offset = PseAlibiComputeOffset<layOutType, hasPse>(pseInfo);
if constexpr (IsSameType<INPUT_T, T>::value) {
if (!pseInfo.align8){
DataCopyIn<INPUT_T, hasPse>(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size,
pseInfo.pseS2Size, alignedSize);
} else {
DataCopyInAlign8<INPUT_T, hasPse>(dstTensor, srcTensor, offset, pseInfo.vec1S1RealSize,
pseInfo.readS2Size, pseInfo.pseS2Size);
}
return;
}
DataCopyIn<INPUT_T, hasPse>(tmpTensor, srcTensor, offset, pseInfo.vec1S1RealSize, pseInfo.readS2Size,
pseInfo.pseS2Size, alignedSize);
if (pseInfo.needCast) {
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize);
}
return;
}
}
template <typename T, bool hasPse>
__aicore__ inline void PseSlopeCopyIn(LocalTensor<T> &dstTensor, LocalTensor<half> &helpTensor,
__gm__ uint8_t *pseSlope, GlobalTensor<half> &alibiGm, PseInfo &pseInfo,
int64_t alignedSize = 16) {
if constexpr (hasPse == true) {
int64_t bOffset = 0;
int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize;
int64_t gOffset = pseInfo.goIdx;
if (pseInfo.pseShapeType == pseSlopeBn) {
bOffset = pseInfo.boIdx * pseInfo.n2G;
}
int64_t offset = bOffset + n2Offset + gOffset;
DataCopyIn<half, hasPse>(helpTensor, alibiGm, 0, pseInfo.vec1S1RealSize,
pseInfo.s2RealSize, pseInfo.pseAlibiBaseS2, alignedSize);
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
if (pseInfo.needCast) {
int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize;
Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize);
AscendC::PipeBarrier<PIPE_V>();
int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset +
pseInfo.loopIdx * pseInfo.vec1S1BaseSize;
int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize;
float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx);
Adds(dstTensor, dstTensor, posShift, computeSize);
AscendC::PipeBarrier<PIPE_V>();
Abs(dstTensor, dstTensor, computeSize);
AscendC::PipeBarrier<PIPE_V>();
float slopes = ((__gm__ T *)pseSlope)[offset] * -1;
if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) {
Sqrt(dstTensor, dstTensor, computeSize);
AscendC::PipeBarrier<PIPE_V>();
}
Muls(dstTensor, dstTensor, slopes, computeSize);
AscendC::PipeBarrier<PIPE_V>();
}
}
}
template <typename T, bool hasPse>
__aicore__ inline void PseSlopeCast(LocalTensor<T> &dstTensor, LocalTensor<half> &helpTensor,
__gm__ uint8_t *pseSlope, PseInfo &pseInfo) {
if constexpr (hasPse == true) {
int64_t bOffset = 0;
int64_t n2Offset = pseInfo.n2oIdx * pseInfo.gSize;
int64_t gOffset = pseInfo.goIdx;
if (pseInfo.pseShapeType == pseSlopeBn) {
bOffset = pseInfo.boIdx * pseInfo.n2G;
}
int64_t offset = bOffset + n2Offset + gOffset;
int64_t computeSize = pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize;
Cast(dstTensor, helpTensor, RoundMode::CAST_NONE, computeSize);
AscendC::PipeBarrier<PIPE_V>();
int64_t s1Offset = pseInfo.s1oIdx * pseInfo.s1BaseSize + pseInfo.vecCoreOffset +
pseInfo.loopIdx * pseInfo.vec1S1BaseSize;
int64_t s2Offset = pseInfo.s2StartIdx + pseInfo.s2LoopCount * pseInfo.s2BaseNratioSize;
float posShift = float(s2Offset + pseInfo.kvStartIdx - s1Offset - pseInfo.qStartIdx);
Adds(dstTensor, dstTensor, posShift, computeSize);
AscendC::PipeBarrier<PIPE_V>();
Abs(dstTensor, dstTensor, computeSize);
AscendC::PipeBarrier<PIPE_V>();
float slopes = ((__gm__ T *)pseSlope)[offset] * -1;
if (pseInfo.pseType == (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) {
Sqrt(dstTensor, dstTensor, computeSize);
AscendC::PipeBarrier<PIPE_V>();
}
Muls(dstTensor, dstTensor, slopes, computeSize);
AscendC::PipeBarrier<PIPE_V>();
}
}
template <typename INPUT_T, typename T, LayOutTypeEnum layOutType, bool hasPse>
__aicore__ inline void PseCopyIn(LocalTensor<T> &dstTensor, LocalTensor<INPUT_T> &tmpTensor,
GlobalTensor<INPUT_T> &srcTensor, PseInfo &pseInfo, int64_t alignedSize = 16)
{
if constexpr (hasPse == true) {
if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) {
return PseAlibiCopyIn<INPUT_T, T, layOutType, hasPse>(dstTensor, tmpTensor, srcTensor, pseInfo, alignedSize);
}
int64_t offset = PseComputeOffset<hasPse>(pseInfo);
int64_t s1Size = pseInfo.pseShapeType == pse1S2 ? (pseInfo.blockCount == 0 ? 1 : pseInfo.blockCount) :
pseInfo.vec1S1RealSize;
if constexpr (IsSameType<INPUT_T, T>::value) {
if (!pseInfo.align8){
DataCopyIn<INPUT_T, hasPse>(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize,
pseInfo.s2Size, alignedSize);
} else {
DataCopyInAlign8<INPUT_T, hasPse>(dstTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size);
}
return;
}
DataCopyIn<INPUT_T, hasPse>(tmpTensor, srcTensor, offset, s1Size, pseInfo.s2RealSize, pseInfo.s2Size,
alignedSize);
if (pseInfo.needCast) {
event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
Cast(dstTensor, tmpTensor, RoundMode::CAST_NONE, s1Size * pseInfo.s2AlignedSize);
}
return;
}
}
template <typename T, bool hasPse>
__aicore__ inline void PseAlibiCompute(LocalTensor<T> &dstTensor, LocalTensor<T> &pseTensor, PseInfo &pseInfo)
{
if constexpr (hasPse == true) {
if (!NeedPseAlibiCompute<hasPse>(pseInfo)) {
return;
}
Add(dstTensor, dstTensor, pseTensor, pseInfo.vec1S1RealSize * pseInfo.pseS2ComputeSize);
return;
}
}
template <typename T, bool hasPse>
__aicore__ inline void PseCompute(LocalTensor<T> &dstTensor, LocalTensor<T> &pseTensor, PseInfo &pseInfo)
{
if constexpr (hasPse == true) {
if (pseInfo.pseEncodeType == pseEncodeALibiS2Full) {
return PseAlibiCompute<T, hasPse>(dstTensor, pseTensor, pseInfo);
}
int64_t computeSize = (pseInfo.pseShapeType == pseS1S2 || pseInfo.pseShapeType == pseSlopeBn ||
pseInfo.pseShapeType == pseSlopeN)
? pseInfo.vec1S1RealSize * pseInfo.s2AlignedSize
: pseInfo.s2AlignedSize;
PseBroadcastAdd<T, hasPse>(pseInfo.vec1S1RealSize, pseInfo.s2AlignedSize, computeSize, pseTensor,
dstTensor, pseInfo.pseShapeType);
return;
}
}
template <bool hasPse>
__aicore__ inline void PseInnerAlibiCreate(GlobalTensor<half> &dstTensor, LocalTensor<half> &helpTensor, PseInfo &pseInfo) {
if constexpr (hasPse == true) {
if (pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_TYPE && pseInfo.pseType != (uint32_t)PseTypeEnum::PSE_INNER_MUL_ADD_SQRT_TYPE) {
return;
}
event_t eventIdMte3ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
event_t eventIdMte3ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
float tmpValue = -1.0;
for (int64_t i = 0; i < pseInfo.pseAlibiBaseS1; i++) {
CreateVecIndex(helpTensor, (half)(i * tmpValue), pseInfo.pseAlibiBaseS2);
SetFlag<HardEvent::V_MTE3>(eventIdVToMte3);
WaitFlag<HardEvent::V_MTE3>(eventIdVToMte3);
DataCopy(dstTensor[i * pseInfo.pseAlibiBaseS2], helpTensor, pseInfo.pseAlibiBaseS2);
SetFlag<HardEvent::MTE3_V>(eventIdMte3ToV);
WaitFlag<HardEvent::MTE3_V>(eventIdMte3ToV);
SetFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
WaitFlag<HardEvent::MTE3_S>(eventIdMte3ToS);
}
}
}
#endif

433
csrc/common/include/kernel/simd.h Normal file
View File

@@ -0,0 +1,433 @@
/**
 * Copyright (c) 2025 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 simd.h
* \brief
*/
#ifndef INCLUDE_SIMD_H
#define INCLUDE_SIMD_H
#ifdef __CCE_KT_TEST__
#define __bf16 bfloat16_t
#endif
#include "hardware.h"
#include "kernel_operator.h"
/////////////////////////////////////////////////////
// vadd
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void add_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
AscendC::LocalTensor<DType> src1,
uint8_t repeat,
uint8_t dstBlockStride,
uint8_t src0BlockStride,
uint8_t src1BlockStride,
uint8_t dstRepeatStride,
uint8_t src0RepeatStride,
uint8_t src1RepeatStride)
{
AscendC::Add<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::BinaryRepeatParams(
dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride));
}
/////////////////////////////////////////////////////
// vadds
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void adds_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
DType scalarValue,
uint8_t repeat,
uint8_t dstBlockStride,
uint8_t srcBlockStride,
uint8_t dstRepeatStride,
uint8_t srcRepeatStride)
{
AscendC::Adds<DType, false>(
dst,
src,
scalarValue,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vcadd
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void cadd_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
uint8_t repeat,
uint16_t dstRepeatStride,
uint16_t srcBlockStride,
uint16_t srcRepeatStride)
{
AscendC::RepeatReduceSum<DType, false>(dst, src, repeat, 0, 0, srcBlockStride, dstRepeatStride, srcRepeatStride);
}
/////////////////////////////////////////////////////
// vbrcb
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void brcb_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
uint16_t dstBlockStride,
uint16_t dstRepeatStride,
uint8_t repeat)
{
AscendC::Brcb(dst, src, repeat, AscendC::BrcbRepeatParams(dstBlockStride, dstRepeatStride));
}
/////////////////////////////////////////////////////
// vcmax
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType, AscendC::ReduceOrder OrderType>
__aicore__ inline void cmax_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
uint8_t repeat,
uint16_t dstRepeatStride,
uint16_t srcBlockStride,
uint16_t srcRepeatStride)
{
#if defined(__DAV_C220_VEC__)
AscendC::WholeReduceMax<DType, false>(
dst, src, (int32_t)0, repeat, dstRepeatStride, srcBlockStride, srcRepeatStride, OrderType);
#else
AscendC::WholeReduceMax<DType, false>(
dst, src, (int32_t)0, repeat, dstRepeatStride, srcBlockStride, srcRepeatStride);
#endif
}
/////////////////////////////////////////////////////
// vconv
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DTypeIn, typename DTypeOut>
__aicore__ inline void conv_v(AscendC::LocalTensor<DTypeOut> dst,
AscendC::LocalTensor<DTypeIn> src,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
if constexpr (std::is_same<DTypeIn, float>::value && std::is_same<DTypeOut, __bf16>::value) {
AscendC::Cast<DTypeOut, DTypeIn, false>(
dst,
src,
AscendC::RoundMode::CAST_RINT,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
} else {
AscendC::Cast<DTypeOut, DTypeIn, false>(
dst,
src,
AscendC::RoundMode::CAST_NONE,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
}
/////////////////////////////////////////////////////
// vconv_f322bf16r
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DTypeIn, typename DTypeOut>
__aicore__ inline void convr_v(AscendC::LocalTensor<DTypeOut> dst,
AscendC::LocalTensor<DTypeIn> src,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Cast<DTypeOut, DTypeIn, false>(
dst,
src,
AscendC::RoundMode::CAST_RINT,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vdiv
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void div_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
AscendC::LocalTensor<DType> src1,
uint8_t repeat,
uint8_t dstBlockStride,
uint8_t src0BlockStride,
uint8_t src1BlockStride,
uint8_t dstRepeatStride,
uint8_t src0RepeatStride,
uint8_t src1RepeatStride)
{
AscendC::Div<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::BinaryRepeatParams(
dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride));
}
/////////////////////////////////////////////////////
// vexp
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void exp_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Exp<DType, false>(
dst,
src,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vmax
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void max_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
AscendC::LocalTensor<DType> src1,
uint8_t repeat,
uint8_t dstBlockStride,
uint8_t src0BlockStride,
uint8_t src1BlockStride,
uint8_t dstRepeatStride,
uint8_t src0RepeatStride,
uint8_t src1RepeatStride)
{
AscendC::Max<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::BinaryRepeatParams(
dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride));
}
/////////////////////////////////////////////////////
// vmul
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void mul_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
AscendC::LocalTensor<DType> src1,
uint8_t repeat,
uint8_t dstBlockStride,
uint8_t src0BlockStride,
uint8_t src1BlockStride,
uint8_t dstRepeatStride,
uint8_t src0RepeatStride,
uint8_t src1RepeatStride)
{
AscendC::Mul<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::BinaryRepeatParams(
dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride));
}
/////////////////////////////////////////////////////
// vmuls
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void muls_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
DType src1,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Muls<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vsub
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void sub_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
AscendC::LocalTensor<DType> src1,
uint8_t repeat,
uint8_t dstBlockStride,
uint8_t src0BlockStride,
uint8_t src1BlockStride,
uint8_t dstRepeatStride,
uint8_t src0RepeatStride,
uint8_t src1RepeatStride)
{
AscendC::Sub<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::BinaryRepeatParams(
dstBlockStride, src0BlockStride, src1BlockStride, dstRepeatStride, src0RepeatStride, src1RepeatStride));
}
/////////////////////////////////////////////////////
// vmaxs
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void maxs_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
DType src1,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Maxs<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vmins
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void mins_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src0,
DType src1,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Mins<DType, false>(
dst,
src0,
src1,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vsqrt
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void sqrt_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Sqrt<DType, false>(
dst,
src,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vln
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void ln_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
uint8_t repeat,
uint16_t dstBlockStride,
uint16_t srcBlockStride,
uint16_t dstRepeatStride,
uint16_t srcRepeatStride)
{
AscendC::Ln<DType, false>(
dst,
src,
(uint64_t)0,
repeat,
AscendC::UnaryRepeatParams(dstBlockStride, srcBlockStride, dstRepeatStride, srcRepeatStride));
}
/////////////////////////////////////////////////////
// vtranspose
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void tranpose_v(AscendC::LocalTensor<DType> dst, AscendC::LocalTensor<DType> src)
{
AscendC::Transpose(dst, src);
}
/////////////////////////////////////////////////////
// vcgmax
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void cgmax_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
const int32_t repeat,
const int32_t dstRepStride,
const int32_t srcBlkStride,
const int32_t srcRepStride)
{
AscendC::BlockReduceMax<DType, false>(dst, src, repeat, 0, dstRepStride, srcBlkStride, srcRepStride);
}
/////////////////////////////////////////////////////
// vcgadd
/////////////////////////////////////////////////////
template <ArchType ArchTag, typename DType>
__aicore__ inline void cgadd_v(AscendC::LocalTensor<DType> dst,
AscendC::LocalTensor<DType> src,
const int32_t repeat,
const int32_t dstRepStride,
const int32_t srcBlkStride,
const int32_t srcRepStride)
{
AscendC::BlockReduceSum<DType, false>(dst, src, repeat, 0, dstRepStride, srcBlkStride, srcRepStride);
}
#endif

159
csrc/common/include/kernel/util.h Normal file
View File

@@ -0,0 +1,159 @@
/**
 * Copyright (c) 2025 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 util.h
* \brief
*/
#ifndef FLASH_ATTENTION_UTIL_H
#define FLASH_ATTENTION_UTIL_H
constexpr int32_t blockBytes = 32;
constexpr int32_t byteBitRatio = 8;
constexpr int64_t prefixAttenMaskDownHeight = 1024;
constexpr static int32_t blockSize = blockBytes / 4; // 4 means sizeof(T)
constexpr static int32_t repeatMaxBytes = 256;
constexpr static int32_t repeatMaxTimes = 255;
constexpr static int32_t repeatMaxSize = repeatMaxBytes / 4; // 4 means sizeof(T)
using AscendC::LocalTensor;
using AscendC::GlobalTensor;
using AscendC::DataFormat;
using AscendC::ShapeInfo;
using AscendC::DataCopyParams;
using AscendC::DataCopyExtParams;
using AscendC::DataCopyPadParams;
using AscendC::DataCopyPadExtParams;
using AscendC::BinaryRepeatParams;
using AscendC::IsSameType;
using AscendC::HardEvent;
using AscendC::SetFlag;
using AscendC::WaitFlag;
enum class LayOutTypeEnum { None = 0, LAYOUT_BSH = 1, LAYOUT_SBH = 2, LAYOUT_BNSD = 3, LAYOUT_TND = 4, LAYOUT_NTD_TND = 5};
namespace math {
template <typename T> __aicore__ inline T Ceil(T a, T b)
{
if (b == 0) {
return 0;
}
return (a + b - 1) / b;
}
template <typename T> __aicore__ inline T Align(T a, T b)
{
if (b == 0) {
return 0;
}
return (a + b - 1) / b * b;
}
}
template <typename T1, typename T2>
__aicore__ inline T1 CeilDiv(T1 a, T2 b)
{
if (b == 0) {
return 0;
}
return (a + b - 1) / b;
}
template <typename T1, typename T2>
__aicore__ inline T1 Max(T1 a, T2 b)
{
return (a > b) ? (a) : (b);
}
template <typename T1, typename T2>
__aicore__ inline T1 Min(T1 a, T2 b)
{
return (a > b) ? (b) : (a);
}
__aicore__ inline void BoolCopyIn(LocalTensor<uint8_t> &dstTensor, GlobalTensor<uint8_t> &srcTensor,
int64_t srcOffset, uint32_t s1Size, uint32_t s2Size, int64_t totalS2Size, int64_t alignedSize = blockBytes)
{
uint32_t alignedS2Size = CeilDiv(s2Size, alignedSize) * alignedSize;
uint32_t shapeArray[] = {s1Size, alignedS2Size};
dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND));
dstTensor.SetSize(s1Size * alignedS2Size);
DataCopyParams dataCopyParams;
dataCopyParams.blockCount = s1Size;
dataCopyParams.dstStride = 0;
if (totalS2Size == blockBytes && alignedSize == 64) { // totalS2Size < 64 && totalS2Size % blockBytes == 0
dataCopyParams.dstStride = 1;
alignedSize = blockBytes;
alignedS2Size = CeilDiv(s2Size, blockBytes) * blockBytes;
}
if (likely(totalS2Size - s2Size <= UINT16_MAX)) {
if (totalS2Size % alignedSize == 0) {
dataCopyParams.blockLen = alignedS2Size / blockBytes;
dataCopyParams.srcStride = (totalS2Size - alignedS2Size) / blockBytes;
DataCopy(dstTensor, srcTensor[srcOffset], dataCopyParams);
} else {
dataCopyParams.blockLen = s2Size;
dataCopyParams.srcStride = totalS2Size - s2Size;
DataCopyPadParams dataCopyPadParams;
dataCopyPadParams.isPad = true;
dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes);
dataCopyPadParams.paddingValue = 1;
DataCopyPad(dstTensor, srcTensor[srcOffset], dataCopyParams, dataCopyPadParams);
}
} else {
DataCopyExtParams extParams;
extParams.blockCount = s1Size;
extParams.dstStride = 0;
extParams.blockLen = s2Size;
extParams.srcStride = totalS2Size - s2Size;
DataCopyPadExtParams<uint8_t> dataCopyPadParams;
dataCopyPadParams.isPad = true;
dataCopyPadParams.rightPadding = Min(alignedS2Size - s2Size, blockBytes);
dataCopyPadParams.paddingValue = 1;
DataCopyPad(dstTensor, srcTensor[srcOffset], extParams, dataCopyPadParams);
}
}
__aicore__ inline void Bit2Int8CopyIn(LocalTensor<uint8_t> &dstTensor, GlobalTensor<uint8_t> &srcTensor,
int64_t srcOffset, uint32_t batchSize, uint32_t s1BaseSize, uint32_t s2BaseSize, int64_t s2TotalSize,
int64_t alignedSize = blockBytes)
{
uint32_t alignedS2Size = CeilDiv(s2BaseSize / byteBitRatio, alignedSize) * alignedSize;
uint32_t shapeArray[] = {batchSize * s1BaseSize, alignedS2Size};
dstTensor.SetShapeInfo(ShapeInfo(2, shapeArray, DataFormat::ND));
dstTensor.SetSize(batchSize * s1BaseSize * alignedS2Size);
DataCopyParams dataCopyParams;
dataCopyParams.blockCount = batchSize * s1BaseSize;
dataCopyParams.blockLen = CeilDiv(s2BaseSize / byteBitRatio, blockBytes);
dataCopyParams.dstStride = 0;
if (s2TotalSize / byteBitRatio % alignedSize == 0 && s2BaseSize / byteBitRatio % alignedSize == 0) {
dataCopyParams.srcStride =
(s2TotalSize / byteBitRatio - dataCopyParams.blockLen * blockBytes) / blockBytes;
DataCopy(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams);
} else {
dataCopyParams.blockLen = CeilDiv(s2BaseSize , byteBitRatio);
dataCopyParams.srcStride = (s2TotalSize - s2BaseSize) / byteBitRatio;
DataCopyPadParams dataCopyPadParams;
dataCopyPadParams.isPad = true;
dataCopyPadParams.rightPadding = 0;
dataCopyPadParams.paddingValue = 0;
DataCopyPad(dstTensor, srcTensor[srcOffset / byteBitRatio], dataCopyParams, dataCopyPadParams);
}
}
__aicore__ inline int32_t Align(int32_t shape)
{
int32_t alignFactor = 16;
int32_t alignedSize = CeilDiv<int32_t, int32_t>(shape, alignFactor) * alignFactor;
return alignedSize;
}
#endif // FLASH_ATTENTION_UTIL_H