from setuptools import setup
from torch.utils import cpp_extension

setup(
    name='btmoe',
    ext_modules=[
        cpp_extension.CUDAExtension(
            'btmoe_cuda', 
            [
                'btmoe/btmoe_cuda.cpp',
                'btmoe/btmoe_cuda_kernel.cu',
                'btmoe/btmoe_cuda_with_zero_kernel.cu'
            ],
            extra_compile_args={
                'cxx': ['-O3', '-DNDEBUG', '-D_GLIBCXX_USE_CXX11_ABI=0'],
                'nvcc': [
                    '-O3',
                    '-DNDEBUG',
                    '-arch=sm_80',
                    '-Xcompiler', '-D_GLIBCXX_USE_CXX11_ABI=0',
                    '--compiler-options', '-fPIC'
                ]
            }
        )
    ],
    cmdclass={'build_ext': cpp_extension.BuildExtension},
    packages=['btmoe'],
    install_requires=['torch']
)
