Commit Graph

1 Commits

Author SHA1 Message Date
1b71e78d13 CUDA support for C++ extensions with setuptools (#5207)
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
2018-02-13 15:02:50 -08:00