Files
pytorch/test/cpp_extensions/setup.py
Roy Li 4c803f4ebd Expose backend extensions to python
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/16582

Reviewed By: gchanan

Differential Revision: D13887539

fbshipit-source-id: 8755babf2e3e849af974655f2f3a91740efe977e
2019-02-01 11:00:18 -08:00

34 lines
1.0 KiB
Python

import sys
import torch.cuda
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
from torch.utils.cpp_extension import CUDA_HOME
CXX_FLAGS = [] if sys.platform == 'win32' else ['-g', '-Werror']
ext_modules = [
CppExtension(
'torch_test_cpp_extension.cpp', ['extension.cpp'],
extra_compile_args=CXX_FLAGS),
CppExtension(
'torch_test_cpp_extension.msnpu', ['msnpu_extension.cpp'],
extra_compile_args=CXX_FLAGS),
]
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': CXX_FLAGS,
'nvcc': ['-O2']})
ext_modules.append(extension)
setup(
name='torch_test_cpp_extension',
packages=['torch_test_cpp_extension'],
ext_modules=ext_modules,
cmdclass={'build_ext': BuildExtension})