You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
24 lines
603 B
24 lines
603 B
from setuptools import setup
|
|
from torch.utils.cpp_extension import CppExtension, CUDAExtension, BuildExtension
|
|
import os
|
|
|
|
this_dir = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
include_dirs = [
|
|
]
|
|
|
|
nvcc_args = [
|
|
'-arch=sm_30',
|
|
'-gencode=arch=compute_30,code=sm_30',
|
|
'-gencode=arch=compute_35,code=sm_35',
|
|
]
|
|
|
|
setup(
|
|
name='ext',
|
|
ext_modules=[
|
|
CppExtension('ext_cpu', ['ext/ext_cpu.cpp']),
|
|
CUDAExtension('ext_cuda', ['ext/ext_cuda.cpp', 'ext/ext_kernel.cu'], extra_compile_args={'cxx': [], 'nvcc': nvcc_args}),
|
|
],
|
|
cmdclass={'build_ext': BuildExtension},
|
|
include_dirs=include_dirs
|
|
)
|
|
|