move files

setup.py

+import os
+import os.path as osp
+import glob
 from setuptools import setup, find_packages
-from sys import argv
+
 import torch
+from torch.utils.cpp_extension import BuildExtension
 from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
 
-TORCH_MAJOR = int(torch.__version__.split('.')[0])
-TORCH_MINOR = int(torch.__version__.split('.')[1])
+WITH_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
+if os.getenv('FORCE_CUDA', '0') == '1':
+    WITH_CUDA = True
+if os.getenv('FORCE_CPU', '0') == '1':
+    WITH_CUDA = False
+
+BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
+
+
+def get_extensions():
+    Extension = CppExtension
+    define_macros = []
+    extra_compile_args = {'cxx': []}
+
+    if WITH_CUDA:
+        Extension = CUDAExtension
+        define_macros += [('WITH_CUDA', None)]
+        nvcc_flags = os.getenv('NVCC_FLAGS', '')
+        nvcc_flags = [] if nvcc_flags == '' else nvcc_flags.split(' ')
+        nvcc_flags += ['-arch=sm_35', '--expt-relaxed-constexpr']
+        extra_compile_args['nvcc'] = nvcc_flags
+
+    extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'csrc')
+    main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
+    extensions = []
+    for main in main_files:
+        name = main.split(os.sep)[-1][:-4]
 
-extra_compile_args = []
-if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
-    extra_compile_args += ['-DVERSION_GE_1_3']
+        sources = [main]
 
-ext_modules = [
-    CppExtension('torch_cluster.graclus_cpu', ['cpu/graclus.cpp'],
-                 extra_compile_args=extra_compile_args),
-    CppExtension('torch_cluster.grid_cpu', ['cpu/grid.cpp']),
-    CppExtension('torch_cluster.fps_cpu', ['cpu/fps.cpp'],
-                 extra_compile_args=extra_compile_args),
-    CppExtension('torch_cluster.rw_cpu', ['cpu/rw.cpp'],
-                 extra_compile_args=extra_compile_args),
-    CppExtension('torch_cluster.sampler_cpu', ['cpu/sampler.cpp'],
-                 extra_compile_args=extra_compile_args),
-]
-cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}
+        path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
+        if osp.exists(path):
+            sources += [path]
 
-GPU = True
-for arg in argv:
-    if arg == '--cpu':
-        GPU = False
-        argv.remove(arg)
+        path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
+        if WITH_CUDA and osp.exists(path):
+            sources += [path]
 
-if CUDA_HOME is not None and GPU:
-    ext_modules += [
-        CUDAExtension('torch_cluster.graclus_cuda',
-                      ['cuda/graclus.cpp', 'cuda/graclus_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-        CUDAExtension('torch_cluster.grid_cuda',
-                      ['cuda/grid.cpp', 'cuda/grid_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-        CUDAExtension('torch_cluster.fps_cuda',
-                      ['cuda/fps.cpp', 'cuda/fps_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-        CUDAExtension('torch_cluster.nearest_cuda',
-                      ['cuda/nearest.cpp', 'cuda/nearest_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-        CUDAExtension('torch_cluster.knn_cuda',
-                      ['cuda/knn.cpp', 'cuda/knn_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-        CUDAExtension('torch_cluster.radius_cuda',
-                      ['cuda/radius.cpp', 'cuda/radius_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-        CUDAExtension('torch_cluster.rw_cuda',
-                      ['cuda/rw.cpp', 'cuda/rw_kernel.cu'],
-                      extra_compile_args=extra_compile_args),
-    ]
+        extension = Extension(
+            'torch_cluster._' + name,
+            sources,
+            include_dirs=[extensions_dir],
+            define_macros=define_macros,
+            extra_compile_args=extra_compile_args,
+        )
+        extensions += [extension]
+
+    return extensions
 
-__version__ = '1.4.5'
-url = 'https://github.com/rusty1s/pytorch_cluster'
 
 install_requires = ['scipy']
 setup_requires = ['pytest-runner']
@@ -63,23 +63,25 @@ tests_require = ['pytest', 'pytest-cov']
 
 setup(
     name='torch_cluster',
-    version=__version__,
-    description=('PyTorch Extension Library of Optimized Graph Cluster '
-                 'Algorithms'),
+    version='1.5.0',
     author='Matthias Fey',
     author_email='matthias.fey@tu-dortmund.de',
-    url=url,
-    download_url='{}/archive/{}.tar.gz'.format(url, __version__),
+    url='https://github.com/rusty1s/pytorch_cluster',
+    description=('PyTorch Extension Library of Optimized Graph Cluster '
+                 'Algorithms'),
     keywords=[
         'pytorch',
         'geometric-deep-learning',
         'graph-neural-networks',
         'cluster-algorithms',
     ],
+    license='MIT',
     install_requires=install_requires,
     setup_requires=setup_requires,
     tests_require=tests_require,
-    ext_modules=ext_modules,
-    cmdclass=cmdclass,
+    ext_modules=get_extensions() if not BUILD_DOCS else [],
+    cmdclass={
+        'build_ext': BuildExtension.with_options(no_python_abi_suffix=True)
+    },
     packages=find_packages(),
 )