/**
* 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.
*/
#include "acl/acl.h"
#include "data_utils.h"
#include "kernel_operator.h"
#ifdef ASCENDC_CPU_DEBUG
#include "cpu_debug_launch.h"
#endif
constexpr uint32_t CUBE_BLOCK = 16;
constexpr uint32_t kLaunchTimes = 20;
constexpr uint32_t M = 256;
constexpr uint32_t K = 64;
constexpr uint32_t N = 256;
constexpr uint32_t singleCoreM = 128;
constexpr uint32_t baseM = 128;
constexpr uint32_t baseK = 64;
constexpr uint32_t baseN = 256;
void ReleaseBuffer(uint8_t*& host, uint8_t*& device)
{
if (device != nullptr) {
(void)aclrtFree(device);
device = nullptr;
}
if (host != nullptr) {
(void)aclrtFreeHost(host);
host = nullptr;
}
}
bool InitInputBuffer(const char* filePath, size_t bufferSize, uint8_t*& host, uint8_t*& device)
{
host = nullptr;
device = nullptr;
aclError aclRet = aclrtMallocHost(reinterpret_cast<void**>(&host), bufferSize);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtMallocHost failed for %s, ret = %d", filePath, static_cast<int32_t>(aclRet));
return false;
}
aclRet = aclrtMalloc(reinterpret_cast<void**>(&device), bufferSize, ACL_MEM_MALLOC_HUGE_FIRST);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtMalloc failed for %s, ret = %d", filePath, static_cast<int32_t>(aclRet));
ReleaseBuffer(host, device);
return false;
}
size_t fileSize = bufferSize;
if (!ReadFile(filePath, fileSize, host, bufferSize)) {
ERROR_LOG("ReadFile failed for %s", filePath);
ReleaseBuffer(host, device);
return false;
}
aclRet = aclrtMemcpy(device, bufferSize, host, bufferSize, ACL_MEMCPY_HOST_TO_DEVICE);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtMemcpy host to device failed for %s, ret = %d", filePath, static_cast<int32_t>(aclRet));
ReleaseBuffer(host, device);
return false;
}
return true;
}
bool InitOutputBuffer(size_t bufferSize, uint8_t*& host, uint8_t*& device)
{
host = nullptr;
device = nullptr;
aclError aclRet = aclrtMallocHost(reinterpret_cast<void**>(&host), bufferSize);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtMallocHost failed for output buffer, ret = %d", static_cast<int32_t>(aclRet));
return false;
}
aclRet = aclrtMalloc(reinterpret_cast<void**>(&device), bufferSize, ACL_MEM_MALLOC_HUGE_FIRST);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtMalloc failed for output buffer, ret = %d", static_cast<int32_t>(aclRet));
ReleaseBuffer(host, device);
return false;
}
return true;
}
template <uint32_t TM, uint32_t TK, uint32_t TN, uint32_t TSingleCoreM, uint32_t TBaseM, uint32_t TBaseK, uint32_t TBaseN>
__global__ __cube__ void mmad_custom(__gm__ uint8_t* a, __gm__ uint8_t* b, __gm__ uint8_t* c)
{
AscendC::GlobalTensor<half> aGM;
AscendC::GlobalTensor<half> bGM;
AscendC::GlobalTensor<half> cGM;
const uint32_t mIterIdx = AscendC::GetBlockIdx() % (TM / TSingleCoreM);
aGM.SetGlobalBuffer((__gm__ half*)a + mIterIdx * TSingleCoreM * TK);
bGM.SetGlobalBuffer((__gm__ half*)b);
cGM.SetGlobalBuffer((__gm__ half*)c + mIterIdx * TSingleCoreM * TN);
AscendC::LocalMemAllocator<AscendC::Hardware::L1> l1Allocator;
AscendC::LocalMemAllocator<AscendC::Hardware::L0A> l0aAllocator;
AscendC::LocalMemAllocator<AscendC::Hardware::L0B> l0bAllocator;
AscendC::LocalMemAllocator<AscendC::Hardware::L0C> l0cAllocator;
AscendC::LocalTensor<half> a1Local = l1Allocator.Alloc<AscendC::TPosition::A1, half>(TBaseM * TBaseK);
AscendC::LocalTensor<half> b1Local = l1Allocator.Alloc<AscendC::TPosition::B1, half>(TBaseK * TBaseN);
AscendC::LocalTensor<half> a2Local = l0aAllocator.Alloc<AscendC::TPosition::A2, half>(TBaseM * TBaseK);
AscendC::LocalTensor<half> b2Local = l0bAllocator.Alloc<AscendC::TPosition::B2, half>(TBaseK * TBaseN);
AscendC::LocalTensor<float> cLocal = l0cAllocator.Alloc<AscendC::TPosition::CO1, float>(TBaseM * TBaseN);
AscendC::DataCopy(a1Local, aGM, AscendC::Nd2NzParams{1, TBaseM, TBaseK, 0, TK, TBaseM, 1, 0});
AscendC::DataCopy(b1Local, bGM, AscendC::Nd2NzParams{1, TBaseK, TBaseN, 0, TN, TBaseK, 1, 0});
AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201)
for (uint32_t i = 0; i < TBaseM / CUBE_BLOCK; ++i) {
AscendC::LoadData(
a2Local[i * TBaseK * CUBE_BLOCK],
a1Local[i * 512 / sizeof(half)],
AscendC::LoadData2DParams{0, TBaseK / CUBE_BLOCK, TBaseM / CUBE_BLOCK, 0, 0, false, 0});
}
for (uint32_t i = 0; i < TBaseK / CUBE_BLOCK; ++i) {
AscendC::LoadData(
b2Local[i * TBaseN * CUBE_BLOCK],
b1Local[i * 512 / sizeof(half)],
AscendC::LoadData2DParams{0, TBaseN / CUBE_BLOCK, TBaseK / CUBE_BLOCK, 0, 0, true, 0});
}
#elif defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)
AscendC::LoadData(
a2Local,
a1Local,
AscendC::LoadData2DParamsV2{
0, 0, TBaseM / CUBE_BLOCK, TBaseK / CUBE_BLOCK, TBaseM / CUBE_BLOCK, TBaseM / CUBE_BLOCK, false, 0});
AscendC::LoadData(
b2Local,
b1Local,
AscendC::LoadData2DParamsV2{
0, 0, TBaseK / CUBE_BLOCK, TBaseN / CUBE_BLOCK, TBaseK / CUBE_BLOCK, TBaseN / CUBE_BLOCK, true, 0});
#endif
AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
AscendC::Mmad(cLocal, a2Local, b2Local, AscendC::MmadParams{TBaseM, TBaseN, TBaseK, 0, false, true});
AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
AscendC::Fixpipe(
cGM,
cLocal,
AscendC::FixpipeParamsV220{TBaseN, TBaseM, TBaseM, TN, false, QuantMode_t::F322F16, 0, 1, 0, 0, 0});
AscendC::PipeBarrier<PIPE_ALL>();
}
int32_t main(int32_t argc, char* argv[])
{
(void)argc;
(void)argv;
const size_t aFileSize = M * K * sizeof(half);
const size_t bFileSize = K * N * sizeof(half);
const size_t cFileSize = M * N * sizeof(half);
const uint32_t numBlocks = M / singleCoreM;
const int32_t deviceId = 0;
aclError aclRet = aclInit(nullptr);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclInit failed, ret = %d", static_cast<int32_t>(aclRet));
return 1;
}
int32_t exitCode = 1;
aclrtStream stream = nullptr;
uint8_t* aHost = nullptr;
uint8_t* aDevice = nullptr;
uint8_t* bHost = nullptr;
uint8_t* bDevice = nullptr;
uint8_t* cHost = nullptr;
uint8_t* cDevice = nullptr;
aclRet = aclrtSetDevice(deviceId);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtSetDevice failed, ret = %d", static_cast<int32_t>(aclRet));
goto finalize_acl;
}
aclRet = aclrtCreateStream(&stream);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtCreateStream failed, ret = %d", static_cast<int32_t>(aclRet));
goto reset_device;
}
if (!InitInputBuffer("./input/x1_gm.bin", aFileSize, aHost, aDevice)) {
goto destroy_stream;
}
if (!InitInputBuffer("./input/x2_gm.bin", bFileSize, bHost, bDevice)) {
goto free_a_buffer;
}
if (!InitOutputBuffer(cFileSize, cHost, cDevice)) {
goto free_b_buffer;
}
for (uint32_t launchIdx = 0; launchIdx < kLaunchTimes; ++launchIdx) {
mmad_custom<M, K, N, singleCoreM, baseM, baseK, baseN><<<numBlocks, 0, stream>>>(aDevice, bDevice, cDevice);
}
aclRet = aclrtSynchronizeStream(stream);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtSynchronizeStream failed, ret = %d", static_cast<int32_t>(aclRet));
goto free_c_buffer;
}
aclRet = aclrtMemcpy(cHost, cFileSize, cDevice, cFileSize, ACL_MEMCPY_DEVICE_TO_HOST);
if (aclRet != ACL_SUCCESS) {
ERROR_LOG("aclrtMemcpy device to host failed, ret = %d", static_cast<int32_t>(aclRet));
goto free_c_buffer;
}
if (!WriteFile("./output/output.bin", cHost, cFileSize)) {
ERROR_LOG("WriteFile failed for ./output/output.bin");
goto free_c_buffer;
}
exitCode = 0;
free_c_buffer:
ReleaseBuffer(cHost, cDevice);
free_b_buffer:
ReleaseBuffer(bHost, bDevice);
free_a_buffer:
ReleaseBuffer(aHost, aDevice);
destroy_stream:
if (stream != nullptr) {
(void)aclrtDestroyStream(stream);
}
reset_device:
(void)aclrtResetDevice(deviceId);
finalize_acl:
(void)aclFinalize();
return exitCode;
}
