/**
* 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.
*/
#define ASCENDC_CUBE_ONLY
#include "acl/acl.h"
#include "data_utils.h"
#include "tiling/tiling_api.h"
#include "kernel_tiling/kernel_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "kernel_operator.h"
#include "lib/matmul_intf.h"
using namespace matmul;
using namespace std;
/**
* @brief Generate matmul tiling.
* @param socVersion: Platform socversion.
* @param tilingBuf data buffer.
*/
void GenerateTiling(platform_ascendc::PlatformAscendC* ascendcPlatform, uint8_t* tilingBuf)
{
int M = 128;
int N = 128;
int K = 256;
matmul_tiling::TPosition leftPosition = matmul_tiling::TPosition::GM;
matmul_tiling::CubeFormat leftFormat = matmul_tiling::CubeFormat::ND;
matmul_tiling::DataType leftDtype = matmul_tiling::DataType::DT_FLOAT16;
bool isTransA = false;
matmul_tiling::TPosition rightPosition = matmul_tiling::TPosition::GM;
matmul_tiling::CubeFormat rightFormat = matmul_tiling::CubeFormat::ND;
matmul_tiling::DataType rightDtype = matmul_tiling::DataType::DT_FLOAT16;
bool isTransB = false;
matmul_tiling::TPosition resultPosition = matmul_tiling::TPosition::GM;
matmul_tiling::CubeFormat resultFormat = matmul_tiling::CubeFormat::ND;
matmul_tiling::DataType resultDtype = matmul_tiling::DataType::DT_FLOAT;
matmul_tiling::TPosition biasPosition = matmul_tiling::TPosition::GM;
matmul_tiling::CubeFormat biasFormat = matmul_tiling::CubeFormat::ND;
matmul_tiling::DataType biasDtype = matmul_tiling::DataType::DT_FLOAT;
bool isBias = true;
int baseM = 128;
int baseN = 128;
optiling::TCubeTiling tilingData;
matmul_tiling::MatmulApiTiling tilingApi(*ascendcPlatform);
tilingApi.SetAType(leftPosition, leftFormat, leftDtype, isTransA);
tilingApi.SetBType(rightPosition, rightFormat, rightDtype, isTransB);
tilingApi.SetCType(resultPosition, resultFormat, resultDtype);
tilingApi.SetBiasType(biasPosition, biasFormat, biasDtype);
tilingApi.SetOrgShape(M, N, K);
tilingApi.SetShape(M, N, K);
tilingApi.SetBias(isBias);
tilingApi.SetTraverse(matmul_tiling::MatrixTraverse::FIRSTM); // Set the matmul travse is FIRSTM.
tilingApi.SetFixSplit(baseM, baseN, -1); // Set the fixed baseM=128, baseN=256.
tilingApi.SetBufferSpace(-1, -1, -1);
int64_t res = tilingApi.GetTiling(tilingData); // Get matmul tiling data.
if (res == -1) {
std::cout << "gen tiling failed" << std::endl;
}
uint32_t tcubeTilingSize = tilingData.GetDataSize();
tilingData.SaveToBuffer(tilingBuf, tcubeTilingSize);
return;
}
__aicore__ inline uint32_t Ceiling(uint32_t a, uint32_t b)
{
return (a + b - 1) / b;
}
/**
* @brief Copy tiling data to TCubeTiling ptr from tiling gm addr.
* @param tiling: TCubeTiling ptr which needs to copy tiling data.
* @param tilingGM: tiling gm addr.
* @retval None
*/
__aicore__ inline void CopyTiling(TCubeTiling* tiling, GM_ADDR tilingGM)
{
uint32_t* ptr = reinterpret_cast<uint32_t*>(tiling);
auto tiling32 = reinterpret_cast<__gm__ uint32_t*>(tilingGM);
for (uint32_t i = 0; i < sizeof(TCubeTiling) / sizeof(uint32_t); i++, ptr++) { *ptr = *(tiling32 + i); }
return;
}
template <typename AType, typename BType, typename CType, typename BiasType>
class MatmulLeakyKernel {
public:
__aicore__ inline MatmulLeakyKernel(){};
__aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c, GM_ADDR workspace,
const TCubeTiling& tiling, AscendC::TPipe* pipe);
__aicore__ inline void Process(AscendC::TPipe* pipe);
__aicore__ inline void CalcOffset(int32_t blockIdx, const TCubeTiling& tiling, int32_t& offsetA, int32_t& offsetB,
int32_t& offsetC, int32_t& offsetBias);
Matmul<MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType>,
MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType>,
MatmulType<AscendC::TPosition::VECIN, CubeFormat::ND, CType>,
MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType>>
matmulObj;
AscendC::GlobalTensor<AType> aGlobal;
AscendC::GlobalTensor<BType> bGlobal;
AscendC::GlobalTensor<CType> cGlobal;
AscendC::GlobalTensor<BiasType> biasGlobal;
TCubeTiling tiling;
};
/**
* @brief Set matmulLeaky input and output gm addr of current core.
* @param a: A matrix gm addr.
* @param b: B matrix gm addr.
* @param bias: Bias gm addr.
* @param c: C matrix gm addr.
* @param workspace: Temporary gm space addr required by matmul calc.
* @param tiling: matmul tiling data.
* @param pipe: Global memory and sync management TPipe object.
* @retval None
*/
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void
MatmulLeakyKernel<AType, BType, CType, BiasType>::Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c, GM_ADDR workspace,
const TCubeTiling& tiling, AscendC::TPipe* pipe)
{
this->tiling = tiling;
aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ AType*>(a), tiling.M * tiling.Ka);
bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ BType*>(b), tiling.Kb * tiling.N);
cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ CType*>(c), tiling.M * tiling.N);
biasGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ BiasType*>(bias), tiling.N);
int32_t offsetA, offsetB, offsetC, offsetBias;
CalcOffset(AscendC::GetBlockIdx(), tiling, offsetA, offsetB, offsetC,
offsetBias); // Calculate the gm offset based on the blockidx.
aGlobal = aGlobal[offsetA];
bGlobal = bGlobal[offsetB];
cGlobal = cGlobal[offsetC];
biasGlobal = biasGlobal[offsetBias];
}
/**
* @brief Main process of matmul calculation
* @param pipe: Global memory and sync management TPipe object.
* @retval None
*/
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulLeakyKernel<AType, BType, CType, BiasType>::Process(AscendC::TPipe* pipe)
{
matmulObj.SetTensorA(aGlobal);
matmulObj.SetTensorB(bGlobal);
matmulObj.SetBias(biasGlobal);
matmulObj.template IterateAll(cGlobal);
matmulObj.End();
AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(3);
}
/**
* @brief Calculate the gm offset based on the blockidx.
* @param blockIdx: Current Core blockidx.
* @param tiling: Matmul tiling data.
* @param offsetA: Gm offset of A matrix.
* @param offsetB: Gm offset of B matrix.
* @param offsetC: Gm offset of C matrix.
* @param offsetBias: Gm offset of Bias matrix.
* @retval None
*/
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void
MatmulLeakyKernel<AType, BType, CType, BiasType>::CalcOffset(int32_t blockIdx, const TCubeTiling& tiling,
int32_t& offsetA, int32_t& offsetB, int32_t& offsetC,
int32_t& offsetBias)
{
auto mSingleBlocks = Ceiling(tiling.M, tiling.singleCoreM);
auto mCoreIndx = blockIdx % mSingleBlocks;
auto nCoreIndx = blockIdx / mSingleBlocks;
offsetA = mCoreIndx * tiling.Ka * tiling.singleCoreM;
offsetB = nCoreIndx * tiling.singleCoreN;
offsetC = mCoreIndx * tiling.N * tiling.singleCoreM + nCoreIndx * tiling.singleCoreN;
offsetBias = nCoreIndx * tiling.singleCoreN;
}
template <typename CType>
class LeakyReluKernel {
public:
__aicore__ inline LeakyReluKernel(){};
__aicore__ inline void Init(GM_ADDR c, const TCubeTiling& tiling, AscendC::TPipe* pipe);
__aicore__ inline void Process(AscendC::TPipe* pipe);
__aicore__ inline void LeakyReluCopyIn(const TCubeTiling& tiling);
__aicore__ inline void LeakyReluCompute(const TCubeTiling& tiling);
__aicore__ inline void LeakyReluCopyOut(const TCubeTiling& tiling);
AscendC::GlobalTensor<CType> cGlobal;
AscendC::LocalTensor<CType> reluInLocal;
AscendC::LocalTensor<CType> reluOutLocal;
TCubeTiling tiling;
AscendC::TQue<AscendC::TPosition::VECIN, 1> reluInQueue_;
AscendC::TQue<AscendC::TPosition::VECOUT, 1> reluOutQueue_;
};
/**
* @brief Set matmulLeaky input and output gm addr of current core.
* @param a: A matrix gm addr.
* @param b: B matrix gm addr.
* @param bias: Bias gm addr.
* @param c: C matrix gm addr.
* @param workspace: Temporary gm space addr required by matmul calc.
* @param tiling: matmul tiling data.
* @param pipe: Global memory and sync management TPipe object.
* @retval None
*/
template <typename CType>
__aicore__ inline void LeakyReluKernel<CType>::Init(GM_ADDR c, const TCubeTiling& tiling, AscendC::TPipe* pipe)
{
this->tiling = tiling;
cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ CType*>(c)
+ AscendC::GetBlockIdx() * tiling.M * tiling.N
/ 2); // c:v = 1:2, split into 2 parts, for vector calculation
pipe->InitBuffer(reluInQueue_, 1,
tiling.singleCoreM * tiling.singleCoreN * sizeof(CType) / 2); // Init input buffer.
pipe->InitBuffer(reluOutQueue_, 1,
tiling.singleCoreM * tiling.singleCoreN * sizeof(CType) / 2); // Init output buffer.
}
template <typename CType>
__aicore__ inline void LeakyReluKernel<CType>::Process(AscendC::TPipe* pipe)
{
AscendC::CrossCoreWaitFlag(3);
LeakyReluCopyIn(tiling);
LeakyReluCompute(tiling);
LeakyReluCopyOut(tiling);
}
template <typename CType>
__aicore__ inline void LeakyReluKernel<CType>::LeakyReluCopyIn(const TCubeTiling& tiling)
{
AscendC::LocalTensor<float> reluInLocal = reluInQueue_.AllocTensor<float>();
AscendC::DataCopy(reluInLocal, cGlobal, tiling.singleCoreM * tiling.singleCoreN / 2);
reluInQueue_.EnQue<float>(reluInLocal);
}
template <typename CType>
__aicore__ inline void LeakyReluKernel<CType>::LeakyReluCompute(const TCubeTiling& tiling)
{
AscendC::LocalTensor<float> reluInLocal = reluInQueue_.DeQue<float>();
AscendC::LocalTensor<float> reluOutLocal = reluOutQueue_.AllocTensor<float>();
AscendC::LeakyRelu(reluOutLocal, reluInLocal, (float)0.001, tiling.singleCoreM * tiling.singleCoreN / 2);
reluOutQueue_.EnQue<float>(reluOutLocal);
reluInQueue_.FreeTensor(reluInLocal);
}
template <typename CType>
__aicore__ inline void LeakyReluKernel<CType>::LeakyReluCopyOut(const TCubeTiling& tiling)
{
AscendC::LocalTensor<float> reluOutLocal = reluOutQueue_.DeQue<float>();
AscendC::DataCopy(cGlobal, reluOutLocal, tiling.singleCoreM * tiling.singleCoreN / 2);
reluOutQueue_.FreeTensor(reluOutLocal);
}
/**
* @brief baremix kernel function entry
* @param a: A matrix gm addr.
* @param b: B matrix gm addr.
* @param bias: Bias gm addr.
* @param c: Out gm addr.
* @param workspace: Temporary gm space addr required by matmul calc.
* @param tilingGm: Tiling data addr.
* @retval None
*/
extern "C" __global__ __aicore__ void baremix_custom(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
GM_ADDR __kfc_workspace__ workspace, GM_ADDR tilingGm)
{
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
AscendC::TPipe pipe;
TCubeTiling tiling;
CopyTiling(&tiling, tilingGm);
if ASCEND_IS_AIC {
MatmulLeakyKernel<half, half, float, float> matmulLeakyKernel;
matmulLeakyKernel.Init(a, b, bias, c, workspace, tiling, &pipe);
REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), matmulLeakyKernel.matmulObj,
&matmulLeakyKernel.tiling); // Initialize the matmul object.
matmulLeakyKernel.Process(&pipe);
}
if ASCEND_IS_AIV {
LeakyReluKernel<float> leakyReluKernel;
leakyReluKernel.Init(c, tiling, &pipe);
leakyReluKernel.Process(&pipe);
}
}
int32_t main(int32_t argc, char* argv[])
{
auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
size_t aFileSize = 32768 * sizeof(int16_t);
size_t bFileSize = 32768 * sizeof(int16_t);
size_t cFileSize = 16384 * sizeof(float);
size_t biasFileSize = 640 * sizeof(float);
size_t tilingFileSize = sizeof(TCubeTiling);
size_t userWorkspaceSize = 0;
size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform->GetLibApiWorkSpaceSize());
size_t workspaceSize = userWorkspaceSize + systemWorkspaceSize;
uint8_t* tilingBuf = (uint8_t*)malloc(tilingFileSize);
GenerateTiling(ascendcPlatform, tilingBuf);
uint32_t numBlocks = 1;
aclInit(nullptr);
int32_t deviceId = 0;
aclrtSetDevice(deviceId);
aclrtStream stream = nullptr;
aclrtCreateStream(&stream);
uint8_t* inputAHost;
uint8_t* inputADevice;
aclrtMallocHost((void**)(&inputAHost), aFileSize);
aclrtMalloc((void**)&inputADevice, aFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
ReadFile("./input/x1_gm.bin", aFileSize, inputAHost, aFileSize);
aclrtMemcpy(inputADevice, aFileSize, inputAHost, aFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
uint8_t* inputBHost;
uint8_t* inputBDevice;
aclrtMallocHost((void**)(&inputBHost), bFileSize);
aclrtMalloc((void**)&inputBDevice, bFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
ReadFile("./input/x2_gm.bin", bFileSize, inputBHost, bFileSize);
aclrtMemcpy(inputBDevice, bFileSize, inputBHost, bFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
uint8_t* outputCHost;
uint8_t* outputCDevice;
aclrtMallocHost((void**)(&outputCHost), cFileSize);
aclrtMalloc((void**)&outputCDevice, cFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
uint8_t* inputBiasHost;
uint8_t* inputBiasDevice;
aclrtMallocHost((void**)(&inputBiasHost), biasFileSize);
aclrtMalloc((void**)&inputBiasDevice, biasFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
ReadFile("./input/bias.bin", biasFileSize, inputBiasHost, biasFileSize);
aclrtMemcpy(inputBiasDevice, biasFileSize, inputBiasHost, biasFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
uint8_t* tilingHost;
uint8_t* tilingDevice;
aclrtMallocHost((void**)(&tilingHost), tilingFileSize);
aclrtMalloc((void**)&tilingDevice, tilingFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
aclrtMemcpy(tilingHost, tilingFileSize, tilingBuf, tilingFileSize, ACL_MEMCPY_HOST_TO_HOST);
aclrtMemcpy(tilingDevice, tilingFileSize, tilingHost, tilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
uint8_t* workspaceDevice;
aclrtMalloc((void**)&workspaceDevice, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
baremix_custom<<<numBlocks, nullptr, stream>>>(inputADevice, inputBDevice, inputBiasDevice, outputCDevice,
workspaceDevice, tilingDevice);
aclrtSynchronizeStream(stream);
aclrtFree(inputADevice);
aclrtFreeHost(inputAHost);
aclrtFree(inputBDevice);
aclrtFreeHost(inputBHost);
aclrtMemcpy(outputCHost, cFileSize, outputCDevice, cFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
WriteFile("./output/output.bin", outputCHost, cFileSize);
aclrtFree(outputCDevice);
aclrtFreeHost(outputCHost);
aclrtFree(inputBiasDevice);
aclrtFreeHost(inputBiasHost);
aclrtFree(tilingDevice);
aclrtFreeHost(tilingHost);
aclrtFree(workspaceDevice);
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
free(tilingBuf);
return 0;
}
