File size: 1,117 Bytes
b3f324b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
# 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
})
|