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