688 lines
31 KiB
C++
688 lines
31 KiB
C++
/**
|
|
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
|
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
|
* CANN Open Software License Agreement Version 2.0 (the "License").
|
|
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
|
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
|
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
|
* See LICENSE in the root of the software repository for the full text of the License.
|
|
*/
|
|
|
|
/*!
|
|
* \file hc_pre_base.h
|
|
* \brief
|
|
*/
|
|
|
|
#ifndef HC_PRE_VECTOR_BASE_H
|
|
#define HC_PRE_VECTOR_BASE_H
|
|
|
|
#include "kernel_operator.h"
|
|
|
|
namespace HcPre {
|
|
using namespace AscendC;
|
|
constexpr int32_t BLOCK_SIZE = 32;
|
|
constexpr int32_t DEFAULT_BLOCK_STRIDE = 1;
|
|
constexpr int32_t DEFAULT_REPEAT_STRIDE = 8;
|
|
constexpr int32_t ONE_REPEAT_BLOCK_NUMS = 8;
|
|
constexpr int32_t REPEAT_SIZE = 256;
|
|
constexpr int32_t MAX_REPEAT_STRIDE = 255;
|
|
constexpr int32_t WORKSPACE_ALIGN_SIZE = 512;
|
|
constexpr uint32_t ONE = 1;
|
|
constexpr uint32_t SHIFT_COEFF = 17;
|
|
constexpr int32_t DOUBLE_BUFFER = 2;
|
|
constexpr int32_t CV_RATIO = 2;
|
|
constexpr uint64_t N_SIZE = 24;
|
|
constexpr uint64_t NUM_TWO = 2;
|
|
constexpr uint64_t SQUARE_SUM_SIZE = 16;
|
|
constexpr uint64_t MASK_PATTERN_BASE_SIZE = 16;
|
|
constexpr uint64_t MASK_PATTERN_REPEAT_SIZE = 8;
|
|
constexpr uint64_t MASK_PATTERN_DIM_SIZE = 16;
|
|
constexpr uint64_t MM_CACHE_LINE_BYTES = 512;
|
|
|
|
__aicore__ inline int32_t CeilDiv(int32_t a, int32_t b)
|
|
{
|
|
if (b == 0) {
|
|
return a;
|
|
}
|
|
return (a + b - 1) / b;
|
|
}
|
|
|
|
__aicore__ inline int32_t CeilAlign(int32_t a, int32_t b)
|
|
{
|
|
return CeilDiv(a, b) * b;
|
|
}
|
|
|
|
template <typename T>
|
|
__aicore__ inline int32_t RoundUp(int32_t num)
|
|
{
|
|
int32_t elemNum = BLOCK_SIZE / sizeof(T);
|
|
return CeilAlign(num, elemNum);
|
|
}
|
|
|
|
__aicore__ inline void SetGatherMaskPattern(const LocalTensor<uint32_t>& maskPattern)
|
|
{
|
|
uint32_t base = ONE;
|
|
for (uint32_t i = 0; i < 16; i++) {
|
|
Duplicate(maskPattern[i * 8], base, 8);
|
|
base = (base << 1);
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
__aicore__ inline void GatherMaskByDiagonal(const LocalTensor<float>& output, const LocalTensor<float>& input,
|
|
const LocalTensor<uint32_t> maskPattern, uint16_t dim0)
|
|
{
|
|
uint32_t totalCount = MASK_PATTERN_DIM_SIZE * MASK_PATTERN_DIM_SIZE;
|
|
uint32_t remainCount = (dim0 % MASK_PATTERN_DIM_SIZE) * MASK_PATTERN_DIM_SIZE;
|
|
uint32_t loopCount = dim0 / MASK_PATTERN_DIM_SIZE;
|
|
uint64_t rsvdCnt = 0;
|
|
for (uint32_t loopIdx = 0; loopIdx < loopCount; loopIdx++) {
|
|
GatherMask(output[loopIdx * MASK_PATTERN_DIM_SIZE],
|
|
input[loopIdx * totalCount], maskPattern, true, totalCount,
|
|
{1, 1, 8, 8}, rsvdCnt);
|
|
}
|
|
GatherMask(output[loopCount * MASK_PATTERN_DIM_SIZE],
|
|
input[loopCount * totalCount], maskPattern, true, remainCount,
|
|
{1, 1, 8, 8}, rsvdCnt);
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
|
|
template <typename T, bool needBrc = true>
|
|
__aicore__ inline void MulABLastDimBrcInline(const LocalTensor<T> &output, const LocalTensor<T> &input0,
|
|
const LocalTensor<T> &input1, const LocalTensor<T> &tmpBuffer,
|
|
const int32_t curRowNum, const int32_t curColNum)
|
|
{
|
|
if constexpr (needBrc) {
|
|
uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS);
|
|
Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE});
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
|
|
uint32_t curColNumAlign = RoundUp<T>(curColNum);
|
|
if (curColNum <= elemInOneBlock) {
|
|
Mul(output, input0, tmpBuffer, curRowNum * curColNumAlign);
|
|
} else {
|
|
int32_t numRepeatPerLine = curColNum / elemInOneRepeat;
|
|
int32_t numRemainPerLine = curColNum % elemInOneRepeat;
|
|
int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock);
|
|
BinaryRepeatParams instrParams;
|
|
|
|
if (numRepeatPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) {
|
|
// 在Col方向开Repeat, 并且Repeat小于255
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Mul(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock],
|
|
elemInOneRepeat, numRepeatPerLine, instrParams);
|
|
}
|
|
} else {
|
|
// 在Row方向开Repeat
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 1;
|
|
for (uint32_t i = 0; i < numRepeatPerLine; i++) {
|
|
Mul(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum,
|
|
instrParams);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (numRemainPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE) {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = 0;
|
|
instrParams.src0RepStride = 0;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Mul(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign],
|
|
input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock],
|
|
numRemainPerLine, 1, instrParams);
|
|
}
|
|
} else {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 0;
|
|
Mul(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer,
|
|
numRemainPerLine, curRowNum, instrParams);
|
|
}
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
template <typename T, bool needBrc = true>
|
|
__aicore__ inline void SubABLastDimBrcInline(const LocalTensor<T> &output, const LocalTensor<T> &input0,
|
|
const LocalTensor<T> &input1, const LocalTensor<T> &tmpBuffer,
|
|
const int32_t curRowNum, const int32_t curColNum)
|
|
{
|
|
if constexpr (needBrc) {
|
|
uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS);
|
|
Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE});
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
|
|
uint32_t curColNumAlign = RoundUp<T>(curColNum);
|
|
if (curColNum <= elemInOneBlock) {
|
|
Sub(output, input0, tmpBuffer, curRowNum * curColNumAlign);
|
|
} else {
|
|
int32_t numRepeatPerLine = curColNum / elemInOneRepeat;
|
|
int32_t numRemainPerLine = curColNum % elemInOneRepeat;
|
|
int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock);
|
|
BinaryRepeatParams instrParams;
|
|
|
|
if (numRepeatPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) {
|
|
// 在Col方向开Repeat, 并且Repeat小于255
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Sub(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock],
|
|
elemInOneRepeat, numRepeatPerLine, instrParams);
|
|
}
|
|
} else {
|
|
// 在Row方向开Repeat
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 1;
|
|
for (uint32_t i = 0; i < numRepeatPerLine; i++) {
|
|
Sub(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum,
|
|
instrParams);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (numRemainPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE) {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = 0;
|
|
instrParams.src0RepStride = 0;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Sub(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign],
|
|
input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock],
|
|
numRemainPerLine, 1, instrParams);
|
|
}
|
|
} else {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 0;
|
|
Sub(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer,
|
|
numRemainPerLine, curRowNum, instrParams);
|
|
}
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
|
|
template <typename T, bool needBrc = true>
|
|
__aicore__ inline void DivABLastDimBrcInline(const LocalTensor<T> &output, const LocalTensor<T> &input0,
|
|
const LocalTensor<T> &input1, const LocalTensor<T> &tmpBuffer,
|
|
const int32_t curRowNum, const int32_t curColNum)
|
|
{
|
|
if constexpr (needBrc) {
|
|
uint32_t repeatTimes = CeilDiv(curRowNum, ONE_REPEAT_BLOCK_NUMS);
|
|
Brcb(tmpBuffer, input1, repeatTimes, {DEFAULT_BLOCK_STRIDE, DEFAULT_REPEAT_STRIDE});
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
|
|
uint32_t curColNumAlign = RoundUp<T>(curColNum);
|
|
if (curColNum <= elemInOneBlock) {
|
|
Div(output, input0, tmpBuffer, curRowNum * curColNumAlign);
|
|
} else {
|
|
int32_t numRepeatPerLine = curColNum / elemInOneRepeat;
|
|
int32_t numRemainPerLine = curColNum % elemInOneRepeat;
|
|
int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock);
|
|
BinaryRepeatParams instrParams;
|
|
|
|
if (numRepeatPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) {
|
|
// 在Col方向开Repeat, 并且Repeat小于255
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Div(output[i * curColNumAlign], input0[i * curColNumAlign], tmpBuffer[i * elemInOneBlock],
|
|
elemInOneRepeat, numRepeatPerLine, instrParams);
|
|
}
|
|
} else {
|
|
// 在Row方向开Repeat
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 1;
|
|
for (uint32_t i = 0; i < numRepeatPerLine; i++) {
|
|
Div(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], tmpBuffer, elemInOneRepeat, curRowNum,
|
|
instrParams);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (numRemainPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE) {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = 0;
|
|
instrParams.src0RepStride = 0;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Div(output[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign],
|
|
input0[numRepeatPerLine * elemInOneRepeat + i * curColNumAlign], tmpBuffer[i * elemInOneBlock],
|
|
numRemainPerLine, 1, instrParams);
|
|
}
|
|
} else {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 0;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 0;
|
|
Div(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], tmpBuffer,
|
|
numRemainPerLine, curRowNum, instrParams);
|
|
}
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
|
|
template <typename T>
|
|
__aicore__ inline void AddBAFirstDimBrcInline(const LocalTensor<T> &output, const LocalTensor<T> &input0,
|
|
const LocalTensor<T> &input1, const int32_t curRowNum,
|
|
const int32_t curColNum)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
|
|
uint32_t curColNumAlign = RoundUp<T>(curColNum);
|
|
int32_t numRepeatPerLine = curColNum / elemInOneRepeat;
|
|
int32_t numRemainPerLine = curColNum % elemInOneRepeat;
|
|
int32_t dstRepStridePerLine = CeilDiv(curColNum, elemInOneBlock);
|
|
BinaryRepeatParams instrParams;
|
|
if (numRepeatPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE || curRowNum < numRepeatPerLine) {
|
|
// 在Col方向开Repeat, 并且Repeat小于255
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Add(output[i * curColNumAlign], input0[i * curColNumAlign], input1, elemInOneRepeat, numRepeatPerLine,
|
|
instrParams);
|
|
}
|
|
} else {
|
|
// 在Row方向开Repeat
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < numRepeatPerLine; i++) {
|
|
Add(output[i * elemInOneRepeat], input0[i * elemInOneRepeat], input1[i * elemInOneRepeat],
|
|
elemInOneRepeat, curRowNum, instrParams);
|
|
}
|
|
}
|
|
}
|
|
|
|
if (numRemainPerLine > 0) {
|
|
if (dstRepStridePerLine > MAX_REPEAT_STRIDE) {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = 0;
|
|
instrParams.src0RepStride = 0;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < curRowNum; i++) {
|
|
Add(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], input1,
|
|
numRemainPerLine, 1, instrParams);
|
|
}
|
|
} else {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = dstRepStridePerLine;
|
|
instrParams.src0RepStride = dstRepStridePerLine;
|
|
instrParams.src1RepStride = 0;
|
|
Add(output[numRepeatPerLine * elemInOneRepeat], input0[numRepeatPerLine * elemInOneRepeat], input1,
|
|
numRemainPerLine, curRowNum, instrParams);
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
template <typename T>
|
|
__aicore__ inline void CalcDenominator(const LocalTensor<T> &output, const LocalTensor<T> &input,
|
|
const uint32_t calCount)
|
|
{
|
|
Muls(output, input, static_cast<T>(-1.0), calCount);
|
|
PipeBarrier<PIPE_V>();
|
|
Exp(output, output, calCount);
|
|
PipeBarrier<PIPE_V>();
|
|
Adds(output, output, static_cast<T>(1.0), calCount);
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
// 暂时不处理repeat超限场景
|
|
template <typename T>
|
|
__aicore__ inline void SigmoidPerf(const LocalTensor<T> &output, const LocalTensor<T> &input,
|
|
const LocalTensor<T> &tmpBuffer, const int64_t calCount)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
Duplicate(tmpBuffer, static_cast<T>(1.0), elemInOneBlock);
|
|
CalcDenominator(output, input, calCount);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
|
|
int32_t numRepeatPerLine = calCount / elemInOneRepeat;
|
|
int32_t numRemainPerLine = calCount % elemInOneRepeat;
|
|
BinaryRepeatParams instrParams;
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 0;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = 0;
|
|
instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE;
|
|
Div(output, tmpBuffer, output, elemInOneRepeat, numRepeatPerLine, instrParams);
|
|
if (numRemainPerLine != 0) {
|
|
Div(output[numRepeatPerLine * elemInOneRepeat], tmpBuffer, output[numRepeatPerLine * elemInOneRepeat],
|
|
numRemainPerLine, 1, instrParams);
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
__aicore__ inline void ProcessPre(const LocalTensor<float> &preLocal, const LocalTensor<float> &mixLocal,
|
|
const LocalTensor<float> &hcBaseLocal, const LocalTensor<float> &rsqrtLocal,
|
|
const LocalTensor<float> &tmpBuffer0, const LocalTensor<float> &tmpBuffer1,
|
|
float scale, float eps, const int32_t curRowNum, const int32_t curColNum)
|
|
{
|
|
int32_t curColNumAlign = RoundUp<float>(curColNum);
|
|
MulABLastDimBrcInline<float, true>(mixLocal, mixLocal, rsqrtLocal, tmpBuffer0, curRowNum, curColNum);
|
|
Muls(mixLocal, mixLocal, scale, curRowNum * curColNumAlign);
|
|
PipeBarrier<PIPE_V>();
|
|
AddBAFirstDimBrcInline<float>(mixLocal, mixLocal, hcBaseLocal, curRowNum, curColNum);
|
|
SigmoidPerf(preLocal, mixLocal, tmpBuffer1, curRowNum * curColNumAlign);
|
|
Adds(preLocal, preLocal, eps, curRowNum * curColNumAlign);
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
__aicore__ inline void ReduceSumARAPerf(const LocalTensor<float> &output, const LocalTensor<float> &input,
|
|
const uint32_t dim0, const uint32_t dim1, const uint32_t dim2)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(float);
|
|
uint32_t dim2Align = RoundUp<float>(dim2);
|
|
|
|
// 拷贝第一个R到output上
|
|
DataCopyParams copyParams;
|
|
copyParams.blockCount = dim0;
|
|
copyParams.blockLen = dim2Align / elemInOneBlock;
|
|
copyParams.srcStride = (dim1 - 1) * (dim2Align / elemInOneBlock);
|
|
copyParams.dstStride = 0;
|
|
DataCopy(output, input, copyParams);
|
|
PipeBarrier<PIPE_V>();
|
|
uint32_t dim2RepeatTimes = dim2 / elemInOneRepeat;
|
|
uint32_t dim2Reminder = dim2 % elemInOneRepeat;
|
|
// 沿着dim2方向开repeat
|
|
BinaryRepeatParams instrParams;
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE;
|
|
for (uint32_t i = 0; i < dim0; i++) {
|
|
for (uint32_t j = 1; j < dim1; j++) {
|
|
Add(output[i * dim2Align], output[i * dim2Align], input[i * dim1 * dim2Align + j * dim2Align],
|
|
elemInOneRepeat, dim2RepeatTimes, instrParams);
|
|
if (dim2Reminder != 0) {
|
|
Add(output[i * dim2Align + dim2RepeatTimes * elemInOneRepeat],
|
|
output[i * dim2Align + dim2RepeatTimes * elemInOneRepeat],
|
|
input[i * dim1 * dim2Align + j * dim2Align + +dim2RepeatTimes * elemInOneRepeat], dim2Reminder, 1,
|
|
instrParams);
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
template <typename T0, typename T1>
|
|
__aicore__ inline void CastTwoDim(const LocalTensor<T0> &output, const LocalTensor<T1> &input, const uint32_t dim0,
|
|
const uint32_t dim1)
|
|
{
|
|
uint32_t dim1AlignT0 = RoundUp<T0>(dim1);
|
|
uint32_t dim1AlignT1 = RoundUp<T1>(dim1);
|
|
if constexpr (IsSameType<T1, bfloat16_t>::value && IsSameType<T0, float>::value) {
|
|
for (uint32_t i = 0; i < dim0; i++) {
|
|
Cast(output[i * dim1AlignT0], input[i * dim1AlignT1], AscendC::RoundMode::CAST_NONE, dim1);
|
|
}
|
|
} else {
|
|
for (uint32_t i = 0; i < dim0; i++) {
|
|
Cast(output[i * dim1AlignT0], input[i * dim1AlignT1], AscendC::RoundMode::CAST_RINT, dim1);
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
|
|
template <typename T>
|
|
__aicore__ void inline ProcessY(const LocalTensor<T> &yLocal, const LocalTensor<T> &xLocal,
|
|
const LocalTensor<float> &mix01Local, const LocalTensor<float> &hcBrcbLocal1,
|
|
const LocalTensor<float> &xCastLocal, const LocalTensor<float> &yCastLocal,
|
|
const uint32_t dim0, const uint32_t dim1, const uint32_t dim2)
|
|
{
|
|
CastTwoDim(xCastLocal, xLocal, dim0 * dim1, dim2);
|
|
MulABLastDimBrcInline<float, true>(xCastLocal, xCastLocal, mix01Local, hcBrcbLocal1, dim0 * dim1, dim2);
|
|
ReduceSumARAPerf(yCastLocal, xCastLocal, dim0, dim1, dim2);
|
|
CastTwoDim(yLocal, yCastLocal, dim0, dim2);
|
|
}
|
|
|
|
|
|
__aicore__ inline void ProcessPost(const LocalTensor<float> &postLocal, const LocalTensor<float> &mixLocal,
|
|
const LocalTensor<float> &hcBaseLocal, const LocalTensor<float> &rsqrtLocal,
|
|
const LocalTensor<float> &tmpBuffer0, const LocalTensor<float> &tmpBuffer1,
|
|
float scale, const int32_t curRowNum, const int32_t curColNum)
|
|
{
|
|
int32_t curColNumAlign = RoundUp<float>(curColNum);
|
|
MulABLastDimBrcInline<float, false>(mixLocal, mixLocal, rsqrtLocal, tmpBuffer0, curRowNum, curColNum);
|
|
Muls(mixLocal, mixLocal, scale, curRowNum * curColNumAlign);
|
|
PipeBarrier<PIPE_V>();
|
|
AddBAFirstDimBrcInline<float>(mixLocal, mixLocal, hcBaseLocal, curRowNum, curColNum);
|
|
SigmoidPerf(postLocal, mixLocal, tmpBuffer1, curRowNum * curColNumAlign);
|
|
Muls(postLocal, postLocal, static_cast<float>(2.0f), curRowNum * curColNumAlign);
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
__aicore__ inline void LastDimReduceMaxPerf(const LocalTensor<float> &output, const LocalTensor<float> &input,
|
|
const uint32_t curRowNum, const uint32_t curColNum)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float);
|
|
WholeReduceMax(output, input, curColNum, curRowNum, 1, 1, CeilDiv(curColNum, elemInOneBlock),
|
|
ReduceOrder::ORDER_ONLY_VALUE);
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
__aicore__ inline void LastDimReduceSumPerf(const LocalTensor<float> &output, const LocalTensor<float> &input,
|
|
const uint32_t curRowNum, const uint32_t curColNum)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(float);
|
|
WholeReduceSum(output, input, curColNum, curRowNum, 1, 1, CeilDiv(curColNum, elemInOneBlock));
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
|
|
// 暂时只支持R轴小于64,既curColNum不能超过64
|
|
__aicore__ inline void SoftmaxFP32Perf(const LocalTensor<float> &output, const LocalTensor<float> &input,
|
|
const LocalTensor<float> &tmpReduceBuffer,
|
|
const LocalTensor<float> tmpBrcbBuffer, const int32_t curRowNum,
|
|
const int32_t curColNum, float eps)
|
|
{
|
|
LastDimReduceMaxPerf(tmpReduceBuffer, input, curRowNum, curColNum);
|
|
SubABLastDimBrcInline<float, true>(output, input, tmpReduceBuffer, tmpBrcbBuffer, curRowNum, curColNum);
|
|
uint32_t curColNumAlign = RoundUp<float>(curColNum);
|
|
Exp(output, output, curRowNum * curColNumAlign);
|
|
PipeBarrier<PIPE_V>();
|
|
LastDimReduceSumPerf(tmpReduceBuffer, output, curRowNum, curColNum);
|
|
DivABLastDimBrcInline<float, true>(output, output, tmpReduceBuffer, tmpBrcbBuffer, curRowNum, curColNum);
|
|
Adds(output, output, eps, curRowNum * curColNumAlign);
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
// (bs, hc_mult, hc_mult) = (bs, hc_mult, hc_mult) + (bs, 1, hc_mult)
|
|
|
|
template <typename T>
|
|
__aicore__ inline void DivABABrcInline(const LocalTensor<T> &output, const LocalTensor<T> &input0,
|
|
const LocalTensor<T> &input1, const uint32_t dim0, const uint32_t dim1,
|
|
const uint32_t dim2)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t elemInOneRepeat = REPEAT_SIZE / sizeof(T);
|
|
uint32_t dim2Align = RoundUp<T>(dim2);
|
|
uint32_t dim2RepeatTimes = dim2 / elemInOneRepeat;
|
|
uint32_t dim2Reminder = dim2 % elemInOneRepeat;
|
|
uint32_t dim2RepeatStride = CeilDiv(dim2, elemInOneBlock);
|
|
// 在dim1方向开repeat
|
|
BinaryRepeatParams instrParams;
|
|
if (dim1 >= dim2RepeatTimes) {
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = dim2RepeatStride;
|
|
instrParams.src0RepStride = dim2RepeatStride;
|
|
instrParams.src1RepStride = 0;
|
|
for (uint32_t i = 0; i < dim0; i++) {
|
|
for (uint32_t j = 0; j < dim2RepeatTimes; j++) {
|
|
Div(output[i * dim1 * dim2Align + j * elemInOneRepeat],
|
|
input0[i * dim1 * dim2Align + j * elemInOneRepeat], input1[i * dim2Align + j * elemInOneRepeat],
|
|
elemInOneRepeat, dim1, instrParams);
|
|
}
|
|
if (dim2Reminder != 0) {
|
|
Div(output[i * dim1 * dim2Align + dim2RepeatTimes * elemInOneRepeat],
|
|
input0[i * dim1 * dim2Align + dim2RepeatTimes * elemInOneRepeat],
|
|
input1[i * dim2Align + dim2RepeatTimes * elemInOneRepeat], dim2Reminder, dim1, instrParams);
|
|
}
|
|
}
|
|
} else {
|
|
// 在dim2方向开repeat
|
|
instrParams.dstBlkStride = 1;
|
|
instrParams.src0BlkStride = 1;
|
|
instrParams.src1BlkStride = 1;
|
|
instrParams.dstRepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src0RepStride = DEFAULT_REPEAT_STRIDE;
|
|
instrParams.src1RepStride = DEFAULT_REPEAT_STRIDE;
|
|
for (uint32_t i = 0; i < dim0; i++) {
|
|
for (uint32_t j = 0; j < dim1; j++) {
|
|
Div(output[i * dim1 * dim2Align + j * dim2Align], input0[i * dim1 * dim2Align + j * dim2Align],
|
|
input1[i * dim2Align], dim2);
|
|
}
|
|
}
|
|
}
|
|
PipeBarrier<PIPE_V>();
|
|
}
|
|
|
|
template <typename T>
|
|
__aicore__ inline void CopyIn(const GlobalTensor<T> &inputGm, const LocalTensor<T> &inputTensor, const uint16_t nBurst,
|
|
const uint32_t copyLen, uint32_t srcStride = 0, uint32_t dstStride = 0)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
DataCopyPadExtParams<T> dataCopyPadExtParams;
|
|
dataCopyPadExtParams.isPad = false;
|
|
dataCopyPadExtParams.leftPadding = 0;
|
|
dataCopyPadExtParams.rightPadding = 0;
|
|
dataCopyPadExtParams.paddingValue = 0;
|
|
|
|
DataCopyExtParams dataCoptExtParams;
|
|
dataCoptExtParams.blockCount = nBurst;
|
|
dataCoptExtParams.blockLen = copyLen * sizeof(T);
|
|
dataCoptExtParams.srcStride = srcStride * sizeof(T);
|
|
dataCoptExtParams.dstStride = dstStride / elemInOneBlock;
|
|
DataCopyPad(inputTensor, inputGm, dataCoptExtParams, dataCopyPadExtParams);
|
|
}
|
|
|
|
// (bs, hc_mult, hc_mult) --> (bs, hc_mult, hc_mult_align)
|
|
template <typename T>
|
|
__aicore__ inline void CopyInWithOuterFor(const GlobalTensor<T> &inputGm, const LocalTensor<T> &inputTensor,
|
|
const uint16_t outerLoop, const uint16_t nBurst, const uint32_t copyLen,
|
|
const uint32_t gmLastDim)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t ubLastDimAlign = RoundUp<T>(copyLen);
|
|
|
|
for (uint16_t i = 0; i < outerLoop; i++) {
|
|
CopyIn(inputGm[i * nBurst * gmLastDim], inputTensor[i * nBurst * ubLastDimAlign], nBurst, copyLen);
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__aicore__ inline void CopyInWithOuterFor(const GlobalTensor<T> &inputGm,
|
|
const LocalTensor<T> &inputTensor, const uint16_t outerLoop,
|
|
const uint16_t nBurst, const uint32_t copyLen, const uint32_t gmFirstDim,
|
|
const uint32_t gmLastDim)
|
|
{
|
|
uint32_t elemInOneBlock = BLOCK_SIZE / sizeof(T);
|
|
uint32_t ubLastDimAlign = RoundUp<T>(copyLen);
|
|
if (outerLoop <= nBurst) {
|
|
for (uint32_t i = 0; i < outerLoop; i++) {
|
|
CopyIn(inputGm[i * gmFirstDim * gmLastDim],
|
|
inputTensor[i * nBurst * ubLastDimAlign], nBurst, copyLen,
|
|
gmLastDim - copyLen);
|
|
}
|
|
} else {
|
|
uint32_t srcStride = (gmLastDim - copyLen) + (gmFirstDim - 1) * gmLastDim;
|
|
uint32_t dstStride = (nBurst - 1) * ubLastDimAlign;
|
|
for (uint32_t i = 0; i < nBurst; i++) {
|
|
CopyIn(inputGm[i * gmLastDim], inputTensor[i * ubLastDimAlign], outerLoop, copyLen, srcStride, dstStride);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
__aicore__ inline void CopyOut(const LocalTensor<T> &outputTensor, const GlobalTensor<T> &outputGm,
|
|
const uint16_t nBurst, const uint32_t copyLen, uint32_t dstStride = 0)
|
|
{
|
|
DataCopyExtParams dataCopyParams;
|
|
dataCopyParams.blockCount = nBurst;
|
|
dataCopyParams.blockLen = copyLen * sizeof(T);
|
|
dataCopyParams.srcStride = 0;
|
|
dataCopyParams.dstStride = dstStride * sizeof(T);
|
|
DataCopyPad(outputGm, outputTensor, dataCopyParams);
|
|
}
|
|
|
|
} // namespace HcPre
|
|
|
|
#endif |