/**
* 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 main.cpp
* \brief Vector Function Add Example (RegBase / VF)
*
* 使用 RegBase（Register Based）编程模型在 Vector Core 上实现向量加法 z = x + y，
* 与 Samples/0_Introduction/vector_add（TQue / MemBase 编程模型）实现相同功能，便于对比。
* 计算在 Vector Function（VF）中完成：数据 Load 进 Vector 寄存器，在寄存器之间完成 Add，
* 结果 Store 回 UB，中间结果不往返 UB。尾块由 Mask 自动处理。
*/
#include <algorithm>
#include <cmath>
#include <iostream>
#include <memory>
#include <random>
#include <tuple>
#include <vector>
#include "acl/acl.h"
#include "kernel_operator.h"
#include "platform/platform_ascendc.h"
// ACL 错误检查宏
#define CHECK_ACL(call) \
do { \
aclError err = (call); \
if (err != ACL_SUCCESS) { \
std::cerr << "ACL error: " << err << " at " << __FILE__ \
<< ":" << __LINE__ << std::endl; \
return 1; \
} \
} while (0)
// 自定义删除器，安全处理空指针
struct AclrtFreeDeleter {
void operator()(void* ptr) const {
if (ptr != nullptr) {
aclrtFree(ptr);
}
}
};
// ----------------------------------------------------------------------------
// 1. Vector Function（VF）：RegBase 编程模型的编程载体
// ----------------------------------------------------------------------------
// VF 用 __simd_vf__ 标记，运行在 Vector Core 上。它直接控制 Vector 寄存器：
// 在函数体内用 AscendC::Reg::* API 完成 Load -> 运算 -> Store，中间结果留在寄存器里。
// 一条指令作用在一整批元素（一个寄存器宽度）上，粒度是 AscendC::VECTOR_REG_WIDTH 字节。
//
// 尾块自适应：total 不一定是寄存器宽度的整数倍。UpdateMask 根据剩余元素数生成掩码，
// 并使剩余数原地递减一个寄存器宽度；最后一轮剩余数不足一个寄存器宽度时，掩码只让
// 前 remain 个元素参与运算，其余 lane 被屏蔽，无需单独处理尾块。
template <typename T>
__simd_vf__ inline void VectorFunctionAdd(
__ubuf__ T* xAddr, __ubuf__ T* yAddr, __ubuf__ T* zAddr, uint32_t total, uint16_t loopNum)
{
constexpr uint32_t vectorLength = AscendC::VECTOR_REG_WIDTH / sizeof(T);
AscendC::Reg::RegTensor<T> xReg, yReg, zReg;
AscendC::Reg::MaskReg mask;
uint32_t remain = total;
for (uint16_t i = 0; i < loopNum; ++i) {
mask = AscendC::Reg::UpdateMask<T>(remain);
AscendC::Reg::LoadAlign<T, AscendC::Reg::LoadDist::DIST_NORM>(
xReg, xAddr + i * vectorLength);
AscendC::Reg::LoadAlign<T, AscendC::Reg::LoadDist::DIST_NORM>(
yReg, yAddr + i * vectorLength);
AscendC::Reg::Add(zReg, xReg, yReg, mask);
AscendC::Reg::StoreAlign<T, AscendC::Reg::StoreDist::DIST_NORM>(
zAddr + i * vectorLength, zReg, mask);
}
}
// ----------------------------------------------------------------------------
// 2. Kernel：GM -> UB -> VF -> UB -> GM
// ----------------------------------------------------------------------------
// 数据从 Global Memory 拷入 UB，由 asc_vf_call 调用 VF 在寄存器上完成计算，
// 结果再从 UB 拷回 Global Memory。多个 block 按 blockLength 切分，每个 block
// 内部按 tileSize 分 tile 处理。
template <typename T>
__global__ __vector__ void vector_function_add_kernel(
__gm__ uint8_t* x, __gm__ uint8_t* y, __gm__ uint8_t* z,
int64_t totalLength, int64_t blockLength, uint32_t tileSize)
{
AscendC::TPipe pipe;
AscendC::TBuf<AscendC::TPosition::VECCALC> xBuf, yBuf, zBuf;
pipe.InitBuffer(xBuf, tileSize);
pipe.InitBuffer(yBuf, tileSize);
pipe.InitBuffer(zBuf, tileSize);
AscendC::LocalTensor<T> xLocal = xBuf.Get<T>();
AscendC::LocalTensor<T> yLocal = yBuf.Get<T>();
AscendC::LocalTensor<T> zLocal = zBuf.Get<T>();
uint32_t blockIdx = AscendC::GetBlockIdx();
AscendC::GlobalTensor<T> xGm, yGm, zGm;
xGm.SetGlobalBuffer((__gm__ T*)x + blockLength * blockIdx);
yGm.SetGlobalBuffer((__gm__ T*)y + blockLength * blockIdx);
zGm.SetGlobalBuffer((__gm__ T*)z + blockLength * blockIdx);
int64_t currentBlockLength = totalLength - blockIdx * blockLength;
currentBlockLength = currentBlockLength > blockLength ? blockLength : currentBlockLength;
if (currentBlockLength <= 0) {
return;
}
constexpr uint32_t vectorLength = AscendC::VECTOR_REG_WIDTH / sizeof(T);
uint32_t elementNumPerTile = tileSize / sizeof(T);
uint32_t tiles = (currentBlockLength + elementNumPerTile - 1) / elementNumPerTile;
for (uint32_t t = 0; t < tiles; ++t) {
int64_t offset = t * elementNumPerTile;
int64_t elems = currentBlockLength - offset;
if (elems > elementNumPerTile) {
elems = elementNumPerTile;
}
uint32_t bytes = static_cast<uint32_t>(elems * sizeof(T));
// GM -> UB：把当前 tile 的 x、y 拷入 UB。尾块不是 32B 整数倍时由 DataCopyPad 补全。
AscendC::DataCopyExtParams copyParams{1, bytes, 0, 0, 0};
AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
AscendC::DataCopyPad(xLocal, xGm[offset], copyParams, padParams);
AscendC::DataCopyPad(yLocal, yGm[offset], copyParams, padParams);
AscendC::PipeBarrier<PIPE_ALL>();
// Compute：调用 VF，Load -> Add -> Store 都在寄存器上完成。
uint16_t loopNum = static_cast<uint16_t>((elems + vectorLength - 1) / vectorLength);
VectorFunctionAdd<T>((__ubuf__ T*)xLocal.GetPhyAddr(), (__ubuf__ T*)yLocal.GetPhyAddr(),
(__ubuf__ T*)zLocal.GetPhyAddr(), static_cast<uint32_t>(elems), loopNum);
AscendC::PipeBarrier<PIPE_ALL>();
// UB -> GM：把计算结果拷回 Global Memory。
AscendC::DataCopyPad(zGm[offset], zLocal, copyParams);
AscendC::PipeBarrier<PIPE_ALL>();
}
}
// ----------------------------------------------------------------------------
// 3. Host 侧：tiling 参数计算
// ----------------------------------------------------------------------------
std::tuple<int64_t, int64_t, int64_t> calc_tiling_params(int64_t totalLength)
{
constexpr static int64_t MIN_ELEMS_PER_CORE = 1024;
auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
uint64_t ubSize;
ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
int64_t coreNum = ascendcPlatform->GetCoreNumAiv();
int64_t numBlocks = std::min(coreNum, (totalLength + MIN_ELEMS_PER_CORE - 1) / MIN_ELEMS_PER_CORE);
numBlocks = std::max(numBlocks, static_cast<int64_t>(1));
int64_t blockLength = (totalLength + numBlocks - 1) / numBlocks;
// x/y/z 三个单缓冲各占一份 UB，tile 大小取 UB 的 1/4，并保证单个 tile 的
// blockLen（uint16）不超过 64KB。
int64_t tileSize = std::min<int64_t>(ubSize / 4, 48 * 1024);
constexpr int64_t ALIGN_SIZE = 32;
tileSize = (tileSize / ALIGN_SIZE) * ALIGN_SIZE;
return std::make_tuple(numBlocks, blockLength, tileSize);
}
int run_vector_add(aclrtStream stream, int64_t numElements)
{
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_real_distribution<float> dist(0.0f, 10.0f);
size_t size = static_cast<size_t>(numElements) * sizeof(float);
// Host 内存
std::vector<float> h_A(numElements);
std::vector<float> h_B(numElements);
std::vector<float> h_C(numElements);
for (int64_t i = 0; i < numElements; ++i) {
h_A[i] = dist(gen);
h_B[i] = dist(gen);
h_C[i] = 0.0f;
}
// Device 内存 - 使用智能指针管理
uint8_t* d_A = nullptr;
uint8_t* d_B = nullptr;
uint8_t* d_C = nullptr;
CHECK_ACL(aclrtMalloc((void**)&d_A, size, ACL_MEM_MALLOC_HUGE_FIRST));
CHECK_ACL(aclrtMalloc((void**)&d_B, size, ACL_MEM_MALLOC_HUGE_FIRST));
CHECK_ACL(aclrtMalloc((void**)&d_C, size, ACL_MEM_MALLOC_HUGE_FIRST));
std::unique_ptr<void, AclrtFreeDeleter> d_A_guard(d_A);
std::unique_ptr<void, AclrtFreeDeleter> d_B_guard(d_B);
std::unique_ptr<void, AclrtFreeDeleter> d_C_guard(d_C);
CHECK_ACL(aclrtMemcpy(d_A, size, h_A.data(), size, ACL_MEMCPY_HOST_TO_DEVICE));
CHECK_ACL(aclrtMemcpy(d_B, size, h_B.data(), size, ACL_MEMCPY_HOST_TO_DEVICE));
// Kernel Call
int64_t numBlocks, blockLength, tileSize;
std::tie(numBlocks, blockLength, tileSize) = calc_tiling_params(numElements);
vector_function_add_kernel<float><<<numBlocks, nullptr, stream>>>(
d_A, d_B, d_C, numElements, blockLength, tileSize);
CHECK_ACL(aclrtSynchronizeStream(stream));
CHECK_ACL(aclrtMemcpy(h_C.data(), size, d_C, size, ACL_MEMCPY_DEVICE_TO_HOST));
CHECK_ACL(aclrtSynchronizeStream(stream));
// 验证结果：浮点加法是精确运算，device 与 host 均按 IEEE 就近舍入，结果逐位一致
bool success = true;
int64_t firstErrorIndex = -1;
for (int64_t i = 0; i < numElements; ++i) {
if (h_C[i] != h_A[i] + h_B[i]) {
success = false;
firstErrorIndex = i;
break;
}
}
if (success) {
std::cout << "Vector function add completed successfully!" << std::endl;
} else {
std::cout << "Vector function add failed at index " << firstErrorIndex << "!" << std::endl;
}
return success ? 0 : 1;
}
int main()
{
CHECK_ACL(aclInit(nullptr));
int32_t deviceId = 0;
CHECK_ACL(aclrtSetDevice(deviceId));
aclrtStream stream = nullptr;
CHECK_ACL(aclrtCreateStream(&stream));
// 409603 不是寄存器宽度（VECTOR_REG_WIDTH / sizeof(float)）的整数倍，
// 用于演示 VF 内部 Mask 对尾块的自适应处理。
int result = run_vector_add(stream, 409603);
CHECK_ACL(aclrtDestroyStream(stream));
CHECK_ACL(aclrtResetDevice(deviceId));
CHECK_ACL(aclFinalize());
return result;
}
