Like this
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;
}
Reference: https://gitcode.com/cann/ops-transformer/blob/master/examples/fast_kernel_launch_example/csrc/add/ascend910b/add.cpp#L161
Like this
Reference: https://gitcode.com/cann/ops-transformer/blob/master/examples/fast_kernel_launch_example/csrc/add/ascend910b/add.cpp#L161