import os
import sys
from setuptools import setup

from torch.utils.cpp_extension import BuildExtension, CUDAExtension
from torch.cuda import is_available

if "--unfused" in sys.argv:
    print("Building unfused version of softsignsgd")
    sys.argv.remove("--unfused")
else:
    cuda_extension = CUDAExtension(
            'fused_softsignsgd', 
            sources=['fused_softsignsgd/pybind_softsignsgd.cpp','./fused_softsignsgd/fused_softsignsgd_kernel.cu', './fused_softsignsgd/multi_tensor_softsignsgd_kernel.cu']
        )

build_cuda_ext = is_available() or os.getenv('FORCE_CUDA', '0') == '1'

setup(
    name='softsignsgd',
    python_requires='>=3.8',
    version='0.0.3',
    install_requires=['torch'],
    py_modules=['softsignsgd'],
    description=(
        'Softsignsgd'
    ),
    author=(
        'qinsh'
    ),
    ext_modules=[cuda_extension] if build_cuda_ext else [],
    cmdclass={'build_ext': BuildExtension} if build_cuda_ext else {},
)
