mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
008ba18c5b
* Also pass torch includes to nvcc build * Export ATen/cuda headers with install * Refactor flags common to C++ and CUDA * Improve tests for C++/CUDA extensions * Export .cuh files under THC * Refactor and clean cpp_extension.py slightly * Include ATen in cuda extension test * Clarifying comment in cuda_extension.cu * Replace cuda_extension.cu with cuda_extension_kernel.cu in setup.py * Copy compile args in C++ extension and add second kernel * Conditionally add -std=c++11 to cuda_flags * Also export cuDNN headers * Add comment about deepcopy
20 lines
718 B
C++
20 lines
718 B
C++
#include <torch/torch.h>
|
|
|
|
// Declare the function from cuda_extension.cu. It will be compiled
|
|
// separately with nvcc and linked with the object file of cuda_extension.cpp
|
|
// into one shared library.
|
|
void sigmoid_add_cuda(const float* x, const float* y, float* output, int size);
|
|
|
|
at::Tensor sigmoid_add(at::Tensor x, at::Tensor y) {
|
|
AT_ASSERT(x.type().is_cuda(), "x must be a CUDA tensor");
|
|
AT_ASSERT(y.type().is_cuda(), "y must be a CUDA tensor");
|
|
auto output = at::zeros_like(x);
|
|
sigmoid_add_cuda(
|
|
x.data<float>(), y.data<float>(), output.data<float>(), output.numel());
|
|
return output;
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("sigmoid_add", &sigmoid_add, "sigmoid(x) + sigmoid(y)");
|
|
}
|