mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
1b71e78d13
This PR adds support for convenient CUDA integration in our C++ extension mechanism. This mainly involved figuring out how to get setuptools to use nvcc for CUDA files and the regular C++ compiler for C++ files. I've added a mixed C++/CUDA test case which works great. I've also added a CUDAExtension and CppExtension function that constructs a setuptools.Extension with "usually the right" arguments, which reduces the required boilerplate to write an extension even more. Especially for CUDA, where library_dir (CUDA_HOME/lib64) and libraries (cudart) have to be specified as well. Next step is to enable this with our "JIT" mechanism. NOTE: I've had to write a small find_cuda_home function to find the CUDA install directory. This logic is kind of a duplicate of tools/setup_helpers/cuda.py, but that's not available in the shipped PyTorch distribution. The function is also fairly short. Let me know if it's fine to duplicate this logic. * CUDA support for C++ extensions with setuptools * Remove printf in CUDA test kernel * Remove -arch flag in test/cpp_extensions/setup.py * Put wrap_compile into BuildExtension * Add guesses for CUDA_HOME directory * export PATH to CUDA location in test.sh * On Python2, sys.platform has the linux version number
20 lines
730 B
C++
20 lines
730 B
C++
#include <torch/torch.h>
|
|
|
|
// Declare the function from cuda_extension_kernel.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_test_cuda_extension, m) {
|
|
m.def("sigmoid_add", &sigmoid_add, "sigmoid(x) + sigmoid(y)");
|
|
}
|