import os 
import shutil
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension

def rmdir(directory):
    if os.path.exists(directory):
        shutil.rmtree(directory)


pwd = os.path.dirname(os.path.abspath(__file__))
rmdir(os.path.join(pwd, 'dist'))
rmdir(os.path.join(pwd, 'build'))
rmdir(os.path.join(pwd, 'sparse_matmul_cuda.egg-info'))


setup(
    name='sparse_matmul_cuda',
    author = 'Qiankun Liu',
    author_email = 'liuqk3@outlook.com',
    ext_modules=[
        CUDAExtension('sparse_matmul_cuda', [
            'src/sparse_matmul_cuda.cpp',
            'src/sparse_matmul_cuda_kernel.cu',],
        ),
    ],
    cmdclass={
        'build_ext': BuildExtension
    })
