1 #include <torch/extension.h>
2
3 // Declare the function from cuda_extension.cu. It will be compiled
4 // separately with nvcc and linked with the object file of cuda_extension.cpp
5 // into one shared library.
6 void sigmoid_add_cuda(const float* x, const float* y, float* output, int size);
7
sigmoid_add(torch::Tensor x,torch::Tensor y)8 torch::Tensor sigmoid_add(torch::Tensor x, torch::Tensor y) {
9 TORCH_CHECK(x.device().is_cuda(), "x must be a CUDA tensor");
10 TORCH_CHECK(y.device().is_cuda(), "y must be a CUDA tensor");
11 auto output = torch::zeros_like(x);
12 sigmoid_add_cuda(
13 x.data_ptr<float>(), y.data_ptr<float>(), output.data_ptr<float>(), output.numel());
14 return output;
15 }
16
PYBIND11_MODULE(TORCH_EXTENSION_NAME,m)17 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
18 m.def("sigmoid_add", &sigmoid_add, "sigmoid(x) + sigmoid(y)");
19 }
20