mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01: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
23 lines
669 B
Python
23 lines
669 B
Python
import torch.cuda
|
|
from setuptools import setup
|
|
from torch.utils.cpp_extension import CppExtension, CUDAExtension
|
|
|
|
ext_modules = [
|
|
CppExtension(
|
|
'torch_test_cpp_extensions', ['extension.cpp'],
|
|
extra_compile_args=['-g']),
|
|
]
|
|
|
|
if torch.cuda.is_available():
|
|
extension = CUDAExtension(
|
|
'torch_test_cuda_extension',
|
|
['cuda_extension.cpp', 'cuda_extension_kernel.cu'],
|
|
extra_compile_args={'cxx': ['-g'],
|
|
'nvcc': ['-O2']})
|
|
ext_modules.append(extension)
|
|
|
|
setup(
|
|
name='torch_test_cpp_extensions',
|
|
ext_modules=ext_modules,
|
|
cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension})
|