/**
* 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"
#define ASCENDC_CUBE_ONLY
#include "lib/matmul_intf.h"
constexpr uint32_t M = 128;
constexpr uint32_t N = 30720;
constexpr uint32_t K = 64;
constexpr bool IS_TRANS_A = false;
constexpr bool IS_TRANS_B = false;
constexpr bool IS_BIAS = false;
constexpr int32_t MAX_M = 10000; // custom matmul kernel support max value of M Dim shape
constexpr int32_t MAX_N = 10000; // custom matmul kernel support max value of N Dim shape
constexpr int32_t MAX_K = 10000; // custom matmul kernel support max value of K Dim shape
constexpr int32_t BASE_M = 128; // BASEM * BASE_K * sizeof(typeC) <=L0A size
constexpr int32_t BASE_N = 256; // BASEN * BASE_K * sizeof(typeB) <=L0B size
constexpr int32_t BASE_K = 64; // BASEM * BASE_N * sizeof(typeC) <=L0C size
constexpr MatmulShapeParams shapeParams = { MAX_M,
MAX_N,
MAX_K,
BASE_M,
BASE_N,
BASE_K };
struct MatmulProblemShape
{
int32_t usedCoreNum;
int32_t m;
int32_t n;
int32_t k;
int32_t sm;
int32_t sn;
int32_t sk;
int32_t isBias;
};
// Calculate single core shapes of per core
MatmulProblemShape GetSingleCoreShape(const int32_t M, const int32_t N, const int32_t K)
{
TCubeTiling tiling;
auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
matmul_tiling::MultiCoreMatmulTiling cubeTiling(*ascendcPlatform);
cubeTiling.SetDim(1);
cubeTiling.SetAType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
cubeTiling.SetBType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
cubeTiling.SetCType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT16);
cubeTiling.SetBiasType(matmul_tiling::TPosition::GM, matmul_tiling::CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
cubeTiling.SetOrgShape(M, N, K);
cubeTiling.SetShape(-1, -1, K);
cubeTiling.EnableBias(IS_BIAS);
if (cubeTiling.GetTiling(tiling) != 0) {
return {};
}
int32_t sm;
int32_t sn;
int32_t sk;
cubeTiling.GetSingleShape(sm, sn, sk);
MatmulProblemShape shapes{tiling.usedCoreNum, M, N, K, sm, sn, sk, IS_BIAS};
return shapes;
}
int32_t GetUsedCoreNum(const MatmulProblemShape& shapes)
{
return shapes.usedCoreNum;
}
constexpr int32_t TILING_DEPTH_PARAM = 8;
constexpr int32_t TILING_STEPK_PARAM = 4;
constexpr int32_t TILING_STEPMN_PARAM = 1;
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline constexpr MatmulApiStaticTiling GetCustomConstantCFG()
{
MatmulConfig mmCFG = GetMMConfig<MatmulConfigMode::CONFIG_MDL>(shapeParams);
// enable unitflag for performance comparison
mmCFG.enUnitFlag = true;
auto constantCFG = AscendC::GetMatmulApiTiling<AType, BType, CType, BiasType>(mmCFG);
constantCFG.depthA1 = TILING_DEPTH_PARAM;
constantCFG.depthB1 = TILING_DEPTH_PARAM;
constantCFG.stepKa = TILING_STEPK_PARAM;
constantCFG.stepKb = TILING_STEPK_PARAM;
constantCFG.stepM = TILING_STEPMN_PARAM;
constantCFG.stepN = TILING_STEPMN_PARAM;
return constantCFG;
}
__aicore__ inline void CopyTiling(MatmulProblemShape *tiling, GM_ADDR tilingGM)
{
int32_t *ptr = reinterpret_cast<int32_t *>(tiling);
auto tiling32 = reinterpret_cast<__gm__ int32_t *>(tilingGM);
for (size_t i = 0; i < sizeof(MatmulProblemShape) / sizeof(int32_t); ++i, ++ptr) {
*ptr = *(tiling32 + i);
}
}
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 tiling: Matmul tiling struct.
* @retval None
*/
__aicore__ inline void Init(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c, GM_ADDR 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);
using A_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, AType>;
using B_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BType>;
using C_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, CType>;
using BIAS_TYPE = AscendC::MatmulType<AscendC::TPosition::GM, CubeFormat::ND, BiasType>;
constexpr static auto CONSTANT_CFG = GetCustomConstantCFG<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE>();
AscendC::Matmul<A_TYPE, B_TYPE, C_TYPE, BIAS_TYPE, CONSTANT_CFG> matmulObj;
MatmulProblemShape shapes;
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(const MatmulProblemShape& param, int32_t& offsetA, int32_t& offsetB,
int32_t& offsetC, int32_t& offsetBias);
AscendC::GlobalTensor<AType> aGlobal;
AscendC::GlobalTensor<BType> bGlobal;
AscendC::GlobalTensor<CType> cGlobal;
AscendC::GlobalTensor<BiasType> biasGlobal;
int32_t mIdx;
int32_t nIdx;
};
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 tiling)
{
CopyTiling(&shapes, tiling);
if (AscendC::GetBlockIdx() >= shapes.usedCoreNum) {
return;
}
aGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ AType*>(a), shapes.m * shapes.k);
bGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ BType*>(b), shapes.n * shapes.k);
cGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ CType*>(c), shapes.m * shapes.n);
biasGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ BiasType*>(bias), shapes.n);
int32_t offsetA = 0;
int32_t offsetB = 0;
int32_t offsetC = 0;
int32_t offsetBias = 0;
aGlobal = aGlobal[offsetA];
bGlobal = bGlobal[offsetB];
cGlobal = cGlobal[offsetC];
biasGlobal = biasGlobal[offsetBias];
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)
{
REGIST_MATMUL_OBJ(pipe, GetSysWorkSpacePtr(), matmulObj, (TCubeTiling*)nullptr);
matmulObj.SetOrgShape(shapes.m, shapes.n, shapes.k);
auto tailM = shapes.m - mIdx * shapes.sm;
tailM = tailM > shapes.sm ? shapes.sm : (tailM > 0 ? tailM : shapes.m);
auto tailN = shapes.n - nIdx * shapes.sn;
tailN = tailN > shapes.sn ? shapes.sn : (tailN > 0 ? tailN : shapes.n);
matmulObj.SetTail(tailM, tailN, shapes.k);
matmulObj.SetTensorA(aGlobal, false);
matmulObj.SetTensorB(bGlobal, false);
if (shapes.isBias) {
matmulObj.SetBias(biasGlobal);
}
matmulObj.IterateAll(cGlobal);
matmulObj.End();
}
template <typename AType, typename BType, typename CType, typename BiasType>
__aicore__ inline void MatmulKernel<AType, BType, CType, BiasType>::CalcOffset(
const MatmulProblemShape& param, int32_t& offsetA, int32_t& offsetB, int32_t& offsetC,
int32_t& offsetBias)
{
auto blockIdx = AscendC::GetBlockIdx();
auto mShape = (param.m + param.sm - 1) / param.sm;
mIdx = blockIdx % mShape;
nIdx = blockIdx / mShape;
offsetA = mIdx * param.k * param.sm;
offsetB = nIdx * param.sn;
offsetC = mIdx * param.n * param.sm + nIdx * param.sn;
offsetBias = nIdx * param.sn;
}
/**
* @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
*/
__cube__ __global__ __aicore__ void matmul_custom(GM_ADDR a, GM_ADDR b, GM_ADDR bias, GM_ADDR c,
GM_ADDR workspace, GM_ADDR tilingGm)
{
if ASCEND_IS_AIV {
return;
}
AscendC::TPipe pipe;
MatmulKernel<half, half, float, float> matmulKernel;
matmulKernel.Init(a, b, bias, c, tilingGm);
matmulKernel.Process(&pipe);
}
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 cFileSize = static_cast<size_t>(M * N) * sizeof(float);
size_t biasFileSize = static_cast<size_t>(N) * sizeof(float);
size_t userWorkspaceSize = 0;
size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform->GetLibApiWorkSpaceSize());
size_t workspaceSize = userWorkspaceSize + systemWorkspaceSize;
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);
const auto shapes = GetSingleCoreShape(M, N, K);
uint32_t coreNum = static_cast<uint32_t>(GetUsedCoreNum(shapes));
uint8_t* tilingHost;
uint8_t* tilingDevice;
size_t tilingFileSize = sizeof(shapes);
aclrtMallocHost((void **)(&tilingHost), tilingFileSize);
aclrtMalloc((void **)&tilingDevice, tilingFileSize,
ACL_MEM_MALLOC_HUGE_FIRST);
aclrtMemcpy(tilingHost, tilingFileSize, &shapes,
tilingFileSize, ACL_MEMCPY_HOST_TO_HOST);
aclrtMemcpy(tilingDevice, tilingFileSize, tilingHost,
tilingFileSize, ACL_MEMCPY_HOST_TO_DEVICE);
matmul_custom<<<coreNum, 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(workspaceDevice);
aclrtFree(tilingDevice);
aclrtFreeHost(tilingHost);
aclrtFree(cDevice);
aclrtFreeHost(cHost);
aclrtDestroyStream(stream);
aclrtResetDevice(deviceId);
aclFinalize();
return 0;
}
