Spaces:
Runtime error
Runtime error
File size: 1,060 Bytes
1bb1365 |
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},
)
|