setup.py

'''
python setup.py sdist bdist_wheel
python -m twine upload dist/*
'''
import setuptools
import glob
import os

from setuptools import find_packages
from setuptools import setup
from torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension, BuildExtension
import sys

requirements = ["torch"]


def get_extensions():
    if CUDA_HOME is None:
        print('CUDA_HOME is None. Install Without CUDA Extension')
        return None
    else:
        print('Install With CUDA Extension')
    this_dir = os.path.dirname(os.path.abspath(__file__))
    extensions_dir = os.path.join(this_dir, 'spikingjelly', 'cext', 'csrc')

    if sys.platform == 'win32':
        # win32 does not support cuSparse
        ext_list = ['neuron']
    else:
        ext_list = ['gemm', 'neuron']
    
    extra_compile_args = {'cxx': ['-g'], 'nvcc': ['-use_fast_math']}
    
    extension = CUDAExtension
    define_macros = [("WITH_CUDA", None)]
    ext_modules = list([
        extension(
            '_C_' + ext_name,
            glob.glob(os.path.join(extensions_dir, ext_name, '*.cpp')) + glob.glob(os.path.join(extensions_dir, ext_name, '*.cu')),
            define_macros=define_macros,
            extra_compile_args=extra_compile_args
        ) for ext_name in ext_list])
    
    return ext_modules

with open("./requirements.txt", "r", encoding="utf-8") as fh:
    install_requires = fh.read()

with open("./README.md", "r", encoding="utf-8") as fh:
    long_description = fh.read()



setup(
    install_requires=install_requires,
    name="spikingjelly",
    version="0.0.0.0.4",
    author="PKU MLG, PCL, and other contributors",
    author_email="fwei@pku.edu.cn, chyq@pku.edu.cn",
    description="A deep learning framework for SNNs built on PyTorch.",
    long_description=long_description,
    long_description_content_type="text/markdown",
    url="https://github.com/fangwei123456/spikingjelly",
    packages=find_packages(),
    classifiers=[
        "Programming Language :: Python :: 3 :: Only",
        "Programming Language :: Python :: 3.6",
        "Programming Language :: Python :: 3.7",
        "Programming Language :: Python :: 3.8",
        "License :: Other/Proprietary License",
        "Operating System :: OS Independent",
    ],
    python_requires='>=3.6',
    ext_modules=get_extensions(),
    cmdclass={
        "build_ext": BuildExtension
    }
)