# 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},
)