mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
b240cc9b87
* Add support for dotted names in CPP Extensions * Modify tests for cpp extensions Test that dotted names work * Py2 fixes * Make run_test cpp_extensions Win-compatible
28 lines
860 B
Python
28 lines
860 B
Python
import torch.cuda
|
|
from setuptools import setup
|
|
from torch.utils.cpp_extension import CppExtension, CUDAExtension
|
|
from torch.utils.cpp_extension import CUDA_HOME
|
|
|
|
ext_modules = [
|
|
CppExtension(
|
|
'torch_test_cpp_extension.cpp', ['extension.cpp'],
|
|
extra_compile_args=['-g']),
|
|
]
|
|
|
|
if torch.cuda.is_available() and CUDA_HOME is not None:
|
|
extension = CUDAExtension(
|
|
'torch_test_cpp_extension.cuda', [
|
|
'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',
|
|
packages=['torch_test_cpp_extension'],
|
|
ext_modules=ext_modules,
|
|
cmdclass={'build_ext': torch.utils.cpp_extension.BuildExtension})
|