/**
* 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 add.asc
* \brief
*/
#include <ATen/Operators.h>
#include <torch/all.h>
#include <torch/library.h>
#include "torch_npu/csrc/core/npu/NPUStream.h"
#include "torch_npu/csrc/framework/OpCommand.h"
#include "kernel_operator.h"
#include "platform/platform_ascendc.h"
#include <type_traits>
namespace ascend_ops {
namespace Add {
// Register the operator's schema
TORCH_LIBRARY_FRAGMENT(EXTENSION_MODULE_NAME, m)
{
m.def("add(Tensor x, Tensor y) -> Tensor");
}
// Meta function implementation of Add
torch::Tensor add_meta(const torch::Tensor &x, const torch::Tensor &y)
{
TORCH_CHECK(x.sizes() == y.sizes(), "The shapes of x and y must be the same.");
auto z = torch::empty_like(x);
return z;
}
// Register the Meta implementation
TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, Meta, m)
{
m.impl("add", add_meta);
}
std::tuple<int64_t, int64_t, int64_t> calc_tiling_params(int64_t totalLength)
{
constexpr static int64_t MIN_ELEMS_PER_CORE = 1024;
constexpr static int64_t PIPELINE_DEPTH = 2;
constexpr static int64_t BUFFER_NUM = 3;
auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance();
uint64_t ubSize;
ascendcPlatform->GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
int64_t coreNum = ascendcPlatform->GetCoreNumAiv();
TORCH_CHECK(coreNum > 0, "coreNum must be positive.");
int64_t numBlocks = std::min(coreNum, (totalLength + MIN_ELEMS_PER_CORE - 1) / MIN_ELEMS_PER_CORE);
int64_t blockLength = (totalLength + numBlocks - 1) / numBlocks;
int64_t tileSize = ubSize / PIPELINE_DEPTH / BUFFER_NUM;
return std::make_tuple(numBlocks, blockLength, tileSize);
}
template <typename T>
__global__ __aicore__ void add_kernel(GM_ADDR x, GM_ADDR y, GM_ADDR z, int64_t totalLength, int64_t blockLength, uint32_t tileSize)
{
constexpr static int64_t PIPELINE_DEPTH = 2;
AscendC::TPipe pipe;
AscendC::GlobalTensor<T> xGm, yGm, zGm;
AscendC::TQue<AscendC::QuePosition::VECIN, PIPELINE_DEPTH> inQueueX;
AscendC::TQue<AscendC::QuePosition::VECIN, PIPELINE_DEPTH> inQueueY;
AscendC::TQue<AscendC::QuePosition::VECOUT, PIPELINE_DEPTH> outQueueZ;
pipe.InitBuffer(inQueueX, PIPELINE_DEPTH, tileSize);
pipe.InitBuffer(inQueueY, PIPELINE_DEPTH, tileSize);
pipe.InitBuffer(outQueueZ, PIPELINE_DEPTH, tileSize);
xGm.SetGlobalBuffer((__gm__ T *)x + blockLength * AscendC::GetBlockIdx());
yGm.SetGlobalBuffer((__gm__ T *)y + blockLength * AscendC::GetBlockIdx());
zGm.SetGlobalBuffer((__gm__ T *)z + blockLength * AscendC::GetBlockIdx());
int64_t currentBlockLength = totalLength - AscendC::GetBlockIdx() * blockLength;
if (currentBlockLength > blockLength) {
currentBlockLength = blockLength;
}
int64_t elementNumPerTile = tileSize / sizeof(T);
int64_t tileNum = currentBlockLength / elementNumPerTile;
int64_t tailTileElementNum = currentBlockLength - tileNum * elementNumPerTile;
for (int64_t i = 0; i < tileNum; ++i) {
int64_t offset = i * elementNumPerTile;
// CopyIn
AscendC::DataCopyExtParams copyParams;
copyParams.blockCount = 1;
copyParams.blockLen = elementNumPerTile * sizeof(T);
copyParams.srcStride = 0;
copyParams.dstStride = 0;
AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
AscendC::LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
AscendC::DataCopyPad(xLocal, xGm[offset], copyParams, padParams);
AscendC::DataCopyPad(yLocal, yGm[offset], copyParams, padParams);
inQueueX.EnQue(xLocal);
inQueueY.EnQue(yLocal);
// Compute
xLocal = inQueueX.DeQue<T>();
yLocal = inQueueY.DeQue<T>();
AscendC::LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
AscendC::Add(zLocal, xLocal, yLocal, elementNumPerTile);
outQueueZ.EnQue(zLocal);
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);
// CopyOut
zLocal = outQueueZ.DeQue<T>();
AscendC::DataCopyPad(zGm[offset], zLocal, copyParams);
outQueueZ.FreeTensor(zLocal);
}
if (tailTileElementNum > 0) {
int64_t offset = tileNum * elementNumPerTile;
// CopyIn
AscendC::DataCopyExtParams copyParams;
copyParams.blockCount = 1;
copyParams.blockLen = tailTileElementNum * sizeof(T);
copyParams.srcStride = 0;
copyParams.dstStride = 0;
AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
AscendC::LocalTensor<T> xLocal = inQueueX.AllocTensor<T>();
AscendC::LocalTensor<T> yLocal = inQueueY.AllocTensor<T>();
AscendC::DataCopyPad(xLocal, xGm[offset], copyParams, padParams);
AscendC::DataCopyPad(yLocal, yGm[offset], copyParams, padParams);
inQueueX.EnQue(xLocal);
inQueueY.EnQue(yLocal);
// Compute
xLocal = inQueueX.DeQue<T>();
yLocal = inQueueY.DeQue<T>();
AscendC::LocalTensor<T> zLocal = outQueueZ.AllocTensor<T>();
AscendC::Add(zLocal, xLocal, yLocal, tailTileElementNum);
outQueueZ.EnQue(zLocal);
inQueueX.FreeTensor(xLocal);
inQueueY.FreeTensor(yLocal);
// CopyOut
zLocal = outQueueZ.DeQue<T>();
AscendC::DataCopyPad(zGm[offset], zLocal, copyParams);
outQueueZ.FreeTensor(zLocal);
}
}
torch::Tensor add_npu(const torch::Tensor &x, const torch::Tensor &y)
{
// OptionalDeviceGuard 确保后续操作在正确的设备上下文执行
// 它会记录当前设备状态，执行完作用域代码后自动恢复
const c10::OptionalDeviceGuard guard(x.device());
auto z = add_meta(x, y);
auto stream = c10_npu::getCurrentNPUStream().stream(false);
int64_t totalLength, numBlocks, blockLength, tileSize;
totalLength = x.numel();
std::tie(numBlocks, blockLength, tileSize) = calc_tiling_params(totalLength);
auto x_ptr = (GM_ADDR)x.data_ptr();
auto y_ptr = (GM_ADDR)y.data_ptr();
auto z_ptr = (GM_ADDR)z.data_ptr();
auto acl_call = [=]() -> int {
AT_DISPATCH_SWITCH(
x.scalar_type(), "add_npu",
AT_DISPATCH_CASE(torch::kFloat32, [&] {
using scalar_t = float;
add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
})
AT_DISPATCH_CASE(torch::kFloat16, [&] {
using scalar_t = half;
add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
})
AT_DISPATCH_CASE(torch::kInt32, [&] {
using scalar_t = int32_t;
add_kernel<scalar_t><<<numBlocks, nullptr, stream>>>(x_ptr, y_ptr, z_ptr, totalLength, blockLength, tileSize);
})
);
return 0;
};
at_npu::native::OpCommand::RunOpApi("Add", acl_call);
return z;
}
// Register the NPU implementation
TORCH_LIBRARY_IMPL(EXTENSION_MODULE_NAME, PrivateUse1, m)
{
m.impl("add", add_npu);
}
} // namespace Add
} // namespace ascend_ops
