35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
# Copyright (C) 2022-present Naver Corporation. All rights reserved.
|
|
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
|
|
|
|
from setuptools import setup
|
|
from torch import cuda
|
|
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
|
|
|
# compile for all possible CUDA architectures
|
|
all_cuda_archs = cuda.get_gencode_flags().replace('compute=','arch=').split()
|
|
# alternatively, you can list cuda archs that you want, eg:
|
|
# all_cuda_archs = [
|
|
# '-gencode', 'arch=compute_70,code=sm_70',
|
|
# '-gencode', 'arch=compute_75,code=sm_75',
|
|
# '-gencode', 'arch=compute_80,code=sm_80',
|
|
# '-gencode', 'arch=compute_86,code=sm_86'
|
|
# ]
|
|
|
|
setup(
|
|
name = 'curope',
|
|
ext_modules = [
|
|
CUDAExtension(
|
|
name='curope',
|
|
sources=[
|
|
"curope.cpp",
|
|
"kernels.cu",
|
|
],
|
|
extra_compile_args = dict(
|
|
nvcc=['-O3','--ptxas-options=-v',"--use_fast_math"]+all_cuda_archs,
|
|
cxx=['-O3'])
|
|
)
|
|
],
|
|
cmdclass = {
|
|
'build_ext': BuildExtension
|
|
})
|