/**
* 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 matmul.asc
* \brief
*/
#include "data_utils.h"
#include "kernel_tiling/kernel_tiling.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"
#include "acl/acl.h"
#include "kernel_operator.h"
#include "lib/matmul_intf.h"
constexpr uint32_t M = 640;
constexpr uint32_t N = 1024;
constexpr uint32_t K = 512;
constexpr bool IS_TRANS_A = false;
constexpr bool IS_TRANS_B = false;
constexpr bool IS_BIAS = true;
/**
* @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 MatmulKernel {
public:
__aicore__ inline MatmulKernel(){};
/**
* @brief Initialization before process.
* @param a: A matrix gm addr.
* @param b: B matrix gm addr.
* @param bias: Bias matrix gm addr.
* @param c: C matrix gm addr.
* @param workspace: Temporary gm space addr required by matmul calc.
* @param tiling: Matmul tiling struct.
* @retval None
*/
__aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c, GM_ADDR workspace,
const TCubeTiling& tiling);
/**
* @brief Process matrix calculation.
* @param pipe: The TPipe object which manages global memory and synchronization.
* @retval None
*/
__aicore__ inline void Process(AscendC::TPipe* pipe);
AscendC::Matmul<AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType, IS_TRANS_A>,
AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType, IS_TRANS_B>,
AscendC::MatmulType<AscendC::TPosition::VECIN, CubeFormat::ND, CType>,
AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType>, CFG_MDL>
matmulObj;
private:
/**
* @brief Calculate the gm offset based on the blockIdx.
* @param blockIdx: Current Core blockidx.
* @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
*/
__aicore__ inline void CalcOffset(int32_t blockIdx, int32_t& offsetA, int32_t& offsetB, int32_t& offsetC,
int32_t& offsetBias);
__aicore__ inline uint32_t CalcDstOffset(uint32_t i);
AscendC::GlobalTensor<AType> aGlobal;
AscendC::GlobalTensor<BType> bGlobal;
AscendC::GlobalTensor<CType> cGlobal;
AscendC::GlobalTensor<BiasType> biasGlobal;
AscendC::GlobalTensor<CType> workspaceGlobal;
AscendC::TQue<AscendC::TPosition::VECIN, 1> cInQueue;
AscendC::TQue<AscendC::TPosition::VECOUT, 1> cOutQueue;
TCubeTiling tiling;
};
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
GM_ADDR workspace, const TCubeTiling& tiling)
{
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);
workspaceGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ CType*>(workspace), tiling.M * tiling.N);
int32_t offsetA = 0;
int32_t offsetB = 0;
int32_t offsetC = 0;
int32_t offsetBias = 0;
CalcOffset(AscendC::GetBlockIdx(), offsetA, offsetB, offsetC, offsetBias);
aGlobal = aGlobal[offsetA];
bGlobal = bGlobal[offsetB];
cGlobal = cGlobal[offsetC];
biasGlobal = biasGlobal[offsetBias];
workspaceGlobal = workspaceGlobal[AscendC::GetBlockIdx() * tiling.singleCoreM * tiling.singleCoreN];
if (GetSysWorkSpacePtr() == nullptr) {
return;
}
}
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::Process(AscendC::TPipe* pipe)
{
matmulObj.SetTensorA(aGlobal, IS_TRANS_A);
matmulObj.SetTensorB(bGlobal, IS_TRANS_B);
if (IS_BIAS) {
matmulObj.SetBias(biasGlobal);
}
matmulObj.SetWorkspace(workspaceGlobal);
matmulObj.template Iterate<false>();
uint32_t baseM = this->tiling.baseM;
uint32_t baseN = this->tiling.baseN;
pipe->InitBuffer(cInQueue, 1, baseM * baseN * sizeof(CType));
pipe->InitBuffer(cOutQueue, 1, baseM * baseN * sizeof(CType));
AscendC::DataCopyParams copyParams = {
(uint16_t)baseM, (uint16_t)(baseN * sizeof(CType) / AscendC::DEFAULT_C0_SIZE), (uint16_t)0,
(uint16_t)((this->tiling.N - baseN) * sizeof(CType) / AscendC::DEFAULT_C0_SIZE)};
uint32_t iterateTimes =
AscendC::Ceil(this->tiling.singleCoreM, baseM) * AscendC::Ceil(this->tiling.singleCoreN, baseN);
for (uint32_t i = 0; i < iterateTimes; ++i) {
// compute
auto cInLocal = cInQueue.AllocTensor<CType>();
matmulObj.template GetTensorC<false>(cInLocal);
cInQueue.EnQue(cInLocal);
// any vector operator
auto src = cInQueue.DeQue<CType>();
auto dst = cOutQueue.AllocTensor<CType>();
DataCopy(dst, src, baseM * baseN);
cOutQueue.EnQue(dst);
cInQueue.FreeTensor(src);
// copy out
auto cOutLocal = cOutQueue.DeQue<CType>();
DataCopy(cGlobal[CalcDstOffset(i)], cOutLocal, copyParams);
cOutQueue.FreeTensor(cOutLocal);
}
matmulObj.End();
}
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::CalcOffset(int32_t blockIdx, int32_t& offsetA,
int32_t& offsetB, int32_t& offsetC,
int32_t& offsetBias)
{
auto mSingleBlocks = AscendC::Ceil(this->tiling.M, this->tiling.singleCoreM);
auto mCoreIndx = blockIdx % mSingleBlocks;
auto nCoreIndx = blockIdx / mSingleBlocks;
offsetA = mCoreIndx * this->tiling.Ka * this->tiling.singleCoreM;
offsetB = nCoreIndx * this->tiling.singleCoreN;
offsetC = mCoreIndx * this->tiling.N * this->tiling.singleCoreM + nCoreIndx * this->tiling.singleCoreN;
offsetBias = nCoreIndx * this->tiling.singleCoreN;
// process with tail block
int32_t tailM = this->tiling.M - mCoreIndx * this->tiling.singleCoreM;
tailM = tailM < this->tiling.singleCoreM ? tailM : this->tiling.singleCoreM;
int32_t tailN = this->tiling.N - nCoreIndx * this->tiling.singleCoreN;
tailN = tailN < this->tiling.singleCoreN ? tailN : this->tiling.singleCoreN;
if (tailM < this->tiling.singleCoreM || tailN < this->tiling.singleCoreN) {
matmulObj.SetTail(tailM, tailN);
}
}
template <typename aType, typename bType, typename CType, typename BiasType>
__aicore__ inline uint32_t MatmulKernel<aType, bType, CType, BiasType>::CalcDstOffset(uint32_t i)
{
uint32_t mIter = 0;
uint32_t nIter = 0;
if (this->tiling.iterateOrder != 1) {
uint32_t mIterTimes = AscendC::Ceil(this->tiling.singleCoreM, this->tiling.baseM);
mIter = i % mIterTimes;
nIter = i / mIterTimes;
} else {
uint32_t nIterTimes = AscendC::Ceil(this->tiling.singleCoreN, this->tiling.baseN);
mIter = i / nIterTimes;
nIter = i % nIterTimes;
}
return (mIter * this->tiling.baseM * this->tiling.N + nIter * this->tiling.baseN);
}
/**
* @brief matmul kernel function entry
* @param a: A matrix gm addr.
* @param b: B matrix gm addr.
* @param bias: bias matrix gm addr.
* @param c: C matrix gm addr.
* @param workspace: Temporary gm space addr required by matmul calc.
* @param tilingGm: Tiling data addr.
* @retval None
*/
__global__ __aicore__ void matmul_custom(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
__kfc_workspace__ GM_ADDR workspace, GM_ADDR tilingGm)
{
TCubeTiling tiling;
CopyTiling(&tiling, tilingGm);
MatmulKernel<half, half, float, float> matmulKernel;
AscendC::TPipe pipe;
REGIST_MATMUL_OBJ(&pipe, GetSysWorkSpacePtr(), matmulKernel.matmulObj, &tiling);
matmulKernel.Init(a, b, bias, c, workspace, tiling);
matmulKernel.Process(&pipe);
}
void GenerateTiling(platform_ascendc::PlatformAscendC* ascendcPlatform, uint8_t* tilingBuf)
{
optiling::TCubeTiling tilingData;
matmul_tiling::MultiCoreMatmulTiling tilingApi(*ascendcPlatform);
tilingApi.SetDim(ascendcPlatform->GetCoreNumAiv());
tilingApi.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16,
IS_TRANS_A);
tilingApi.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16,
IS_TRANS_B);
tilingApi.SetCType(matmul_tiling::TPosition::VECIN, matmul_tiling::CubeFormat::ND,
matmul_tiling::DataType::DT_FLOAT);
tilingApi.SetBiasType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND,
matmul_tiling::DataType::DT_FLOAT);
tilingApi.SetOrgShape(M, N, K);
tilingApi.SetShape(M, N, K);
tilingApi.EnableBias(IS_BIAS);
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);
}
int32_t main(int32_t argc, char* argv[])
{
auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
size_t aFileSize = static_cast<size_t>(M * K) * sizeof(uint16_t); // uint16_t represent half
size_t bFileSize = static_cast<size_t>(K * N) * sizeof(uint16_t); // uint16_t represent half
size_t biasFileSize = static_cast<size_t>(sizeof(float) * N);
size_t cFileSize = static_cast<size_t>(M * N) * sizeof(float);
size_t userWorkspaceSize = static_cast<size_t>(sizeof(float) * M * N);
size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform->GetLibApiWorkSpaceSize());
size_t workspaceSize = userWorkspaceSize + systemWorkspaceSize;
// matmul TCubeTiling
size_t tilingFileSize = sizeof(TCubeTiling);
uint8_t* tilingBuf = (uint8_t*)malloc(tilingFileSize);
GenerateTiling(ascendcPlatform, tilingBuf);
uint32_t numBlocks = reinterpret_cast<TCubeTiling*>(tilingBuf)->usedCoreNum / 2; // AIC:AIV = 1:2
int32_t deviceId = 0;
aclrtStream stream = nullptr;
aclrtContext context;
aclInit(nullptr);
aclrtSetDevice(deviceId);
aclrtCreateContext(&context, deviceId);
aclrtCreateStream(&stream);
uint8_t* aHost;
uint8_t* aDevice;
aclrtMallocHost((void**)(&aHost), aFileSize);
aclrtMalloc((void**)&aDevice, aFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
ReadFile("./input/x1_gm.bin", aFileSize, aHost, aFileSize);
aclrtMemcpy(aDevice, aFileSize, aHost, aFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
uint8_t* bHost;
uint8_t* bDevice;
aclrtMallocHost((void**)(&bHost), bFileSize);
aclrtMalloc((void**)&bDevice, bFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
ReadFile("./input/x2_gm.bin", bFileSize, bHost, bFileSize);
aclrtMemcpy(bDevice, bFileSize, bHost, bFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
uint8_t* biasHost;
uint8_t* biasDevice;
if (IS_BIAS) {
aclrtMallocHost((void**)(&biasHost), biasFileSize);
aclrtMalloc((void**)&biasDevice, biasFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
ReadFile("./input/bias_gm.bin", biasFileSize, biasHost, biasFileSize);
aclrtMemcpy(biasDevice, biasFileSize, biasHost, biasFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
}
uint8_t* cHost;
uint8_t* cDevice;
aclrtMallocHost((void**)(&cHost), cFileSize);
aclrtMalloc((void**)&cDevice, cFileSize, ACL_MEM_MALLOC_HUGE_FIRST);
uint8_t* workspaceDevice;
aclrtMalloc((void**)&workspaceDevice, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
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);
matmul_custom<<<numBlocks, nullptr, stream>>>(aDevice, bDevice, biasDevice, cDevice, workspaceDevice, tilingDevice);
aclrtSynchronizeStream(stream);
aclrtMemcpy(cHost, cFileSize, cDevice, cFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
WriteFile("./output/output.bin", cHost, cFileSize);
aclrtFree(aDevice);
aclrtFreeHost(aHost);
aclrtFree(bDevice);
aclrtFreeHost(bHost);
if (IS_BIAS) {
aclrtFree(biasDevice);
aclrtFreeHost(biasHost);
}
aclrtFree(cDevice);
aclrtFreeHost(cHost);
aclrtFree(workspaceDevice);
aclrtFree(tilingDevice);
aclrtFreeHost(tilingHost);
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
free(tilingBuf);
return 0;
}
