build separate cpu and cuda images

rusty1s · rusty1s · commit 6ab5c3322cb7 · 2021-03-03T07:51:03.000+01:00
diff --git a/.travis.yml b/.travis.yml
@@ -112,15 +112,16 @@ install:
   - conda install pytorch=${TORCH_VERSION} ${TOOLKIT} -c pytorch --yes
   - source script/torch.sh
   - pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-${TORCH_VERSION}+${IDX}.html
-  - pip install flake8 codecov
+  - pip install flake8
+  - pip install codecov
   - pip install scipy==1.4.1
-  - source script/install.sh
+  - travis_wait 30 pip install -e .
 script:
   - flake8 .
   - python setup.py test
 after_success:
-  - python setup.py bdist_wheel --dist-dir=dist/torch-${TORCH_VERSION}
-  - python script/rename_wheel.py ${IDX}
+  - python setup.py bdist_wheel --dist-dir=dist
+  - ls -lah dist/
   - codecov
 deploy:
   provider: s3
@@ -129,8 +130,8 @@ deploy:
   access_key_id: ${S3_ACCESS_KEY}
   secret_access_key: ${S3_SECRET_ACCESS_KEY}
   bucket: pytorch-geometric.com
-  local_dir: dist/torch-${TORCH_VERSION}
-  upload_dir: whl/torch-${TORCH_VERSION}
+  local_dir: dist
+  upload_dir: whl/torch-${TORCH_VERSION}+${IDX}
   acl: public_read
   on:
     all_branches: true
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.0)
 project(torchsparse)
 set(CMAKE_CXX_STANDARD 14)
-set(TORCHSPARSE_VERSION 0.6.8)
+set(TORCHSPARSE_VERSION 0.6.9)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
 
diff --git a/csrc/convert.cpp b/csrc/convert.cpp
@@ -8,7 +8,11 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__convert(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__convert_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; }
+#endif
 #endif
 
 torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
diff --git a/csrc/diag.cpp b/csrc/diag.cpp
@@ -8,7 +8,11 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__diag(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__diag_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__diag_cpu(void) { return NULL; }
+#endif
 #endif
 
 torch::Tensor non_diag_mask(torch::Tensor row, torch::Tensor col, int64_t M,
diff --git a/csrc/metis.cpp b/csrc/metis.cpp
@@ -4,7 +4,11 @@
 #include "cpu/metis_cpu.h"
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__metis(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__metis_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__metis_cpu(void) { return NULL; }
+#endif
 #endif
 
 torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
diff --git a/csrc/relabel.cpp b/csrc/relabel.cpp
@@ -4,7 +4,11 @@
 #include "cpu/relabel_cpu.h"
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__relabel(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__relablel_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__relabel_cpu(void) { return NULL; }
+#endif
 #endif
 
 std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
diff --git a/csrc/rw.cpp b/csrc/rw.cpp
@@ -8,7 +8,11 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__rw(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__rw_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__rw_cpu(void) { return NULL; }
+#endif
 #endif
 
 torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
diff --git a/csrc/saint.cpp b/csrc/saint.cpp
@@ -4,7 +4,11 @@
 #include "cpu/saint_cpu.h"
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__saint(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__saint_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__saint_cpu(void) { return NULL; }
+#endif
 #endif
 
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor>
diff --git a/csrc/sample.cpp b/csrc/sample.cpp
@@ -4,7 +4,11 @@
 #include "cpu/sample_cpu.h"
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__sample(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__sample_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__sample_cpu(void) { return NULL; }
+#endif
 #endif
 
 std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
diff --git a/csrc/spmm.cpp b/csrc/spmm.cpp
@@ -8,7 +8,11 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__spmm(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__spmm_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__spmm_cpu(void) { return NULL; }
+#endif
 #endif
 
 std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
diff --git a/csrc/spspmm.cpp b/csrc/spspmm.cpp
@@ -8,7 +8,11 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__spspmm(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__spspmm_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__spspmm_cpu(void) { return NULL; }
+#endif
 #endif
 
 std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
diff --git a/csrc/version.cpp b/csrc/version.cpp
@@ -6,7 +6,11 @@
 #endif
 
 #ifdef _WIN32
-PyMODINIT_FUNC PyInit__version(void) { return NULL; }
+#ifdef WITH_CUDA
+PyMODINIT_FUNC PyInit__version_cuda(void) { return NULL; }
+#else
+PyMODINIT_FUNC PyInit__version_cpu(void) { return NULL; }
+#endif
 #endif
 
 int64_t cuda_version() {
diff --git a/script/cuda.sh b/script/cuda.sh
@@ -69,7 +69,7 @@ if [ "${TRAVIS_OS_NAME}" = "osx" ] && [ "$IDX" = "cpu" ]; then
 fi
 
 if [ "${IDX}" = "cpu" ]; then
-  export FORCE_CPU=1
+  export FORCE_ONLY_CPU=1
 else
   export FORCE_CUDA=1
 fi
diff --git a/script/rename_wheel.py b/script/rename_wheel.py
diff --git a/setup.py b/setup.py
@@ -1,7 +1,8 @@
 import os
-import os.path as osp
 import sys
 import glob
+import os.path as osp
+from itertools import product
 from setuptools import setup, find_packages
 
 import torch
@@ -10,10 +11,13 @@
 from torch.utils.cpp_extension import CppExtension, CUDAExtension, CUDA_HOME
 
 WITH_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
+suffices = ['cpu', 'cuda'] if WITH_CUDA else ['cpu']
 if os.getenv('FORCE_CUDA', '0') == '1':
-    WITH_CUDA = True
-if os.getenv('FORCE_CPU', '0') == '1':
-    WITH_CUDA = False
+    suffices = ['cuda', 'cpu']
+if os.getenv('FORCE_ONLY_CUDA', '0') == '1':
+    suffices = ['cuda']
+if os.getenv('FORCE_ONLY_CPU', '0') == '1':
+    suffices = ['cpu']
 
 BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
 
@@ -22,63 +26,63 @@
 
 
 def get_extensions():
-    Extension = CppExtension
-    define_macros = []
-    libraries = []
-    if WITH_METIS:
-        define_macros += [('WITH_METIS', None)]
-        libraries += ['metis']
-    if WITH_MTMETIS:
-        define_macros += [('WITH_MTMETIS', None)]
-        define_macros += [('MTMETIS_64BIT_VERTICES', None)]
-        define_macros += [('MTMETIS_64BIT_EDGES', None)]
-        define_macros += [('MTMETIS_64BIT_WEIGHTS', None)]
-        define_macros += [('MTMETIS_64BIT_PARTITIONS', None)]
-        libraries += ['mtmetis', 'wildriver']
-    extra_compile_args = {'cxx': ['-O2']}
-    extra_link_args = ['-s']
-
-    info = parallel_info()
-    if 'parallel backend: OpenMP' in info and 'OpenMP not found' not in info:
-        extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
-        if sys.platform == 'win32':
-            extra_compile_args['cxx'] += ['/openmp']
-        else:
-            extra_compile_args['cxx'] += ['-fopenmp']
-    else:
-        print('Compiling without OpenMP...')
-
-    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', '-O2']
-        extra_compile_args['nvcc'] = nvcc_flags
-
-        if sys.platform == 'win32':
-            extra_link_args += ['cusparse.lib']
-        else:
-            extra_link_args += ['-lcusparse', '-l', 'cusparse']
+    extensions = []
 
     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]
 
+    for main, suffix in product(main_files, suffices):
+        define_macros = []
+        libraries = []
+        if WITH_METIS:
+            define_macros += [('WITH_METIS', None)]
+            libraries += ['metis']
+        if WITH_MTMETIS:
+            define_macros += [('WITH_MTMETIS', None)]
+            define_macros += [('MTMETIS_64BIT_VERTICES', None)]
+            define_macros += [('MTMETIS_64BIT_EDGES', None)]
+            define_macros += [('MTMETIS_64BIT_WEIGHTS', None)]
+            define_macros += [('MTMETIS_64BIT_PARTITIONS', None)]
+            libraries += ['mtmetis', 'wildriver']
+        extra_compile_args = {'cxx': ['-O2']}
+        extra_link_args = ['-s']
+
+        info = parallel_info()
+        if 'backend: OpenMP' in info and 'OpenMP not found' not in info:
+            extra_compile_args['cxx'] += ['-DAT_PARALLEL_OPENMP']
+            if sys.platform == 'win32':
+                extra_compile_args['cxx'] += ['/openmp']
+            else:
+                extra_compile_args['cxx'] += ['-fopenmp']
+        else:
+            print('Compiling without OpenMP...')
+
+        if suffix == 'cuda':
+            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', '-O2']
+            extra_compile_args['nvcc'] = nvcc_flags
+
+            if sys.platform == 'win32':
+                extra_link_args += ['cusparse.lib']
+            else:
+                extra_link_args += ['-lcusparse', '-l', 'cusparse']
+
+        name = main.split(os.sep)[-1][:-4]
         sources = [main]
 
         path = osp.join(extensions_dir, 'cpu', f'{name}_cpu.cpp')
         if osp.exists(path):
             sources += [path]
 
         path = osp.join(extensions_dir, 'cuda', f'{name}_cuda.cu')
-        if WITH_CUDA and osp.exists(path):
+        if suffix == 'cuda' and osp.exists(path):
             sources += [path]
 
+        Extension = CppExtension if suffix == 'cpu' else CUDAExtension
         extension = Extension(
-            'torch_sparse._' + name,
+            f'torch_sparse._{name}_{suffix}',
             sources,
             include_dirs=[extensions_dir],
             define_macros=define_macros,
@@ -97,7 +101,7 @@ def get_extensions():
 
 setup(
     name='torch_sparse',
-    version='0.6.8',
+    version='0.6.9',
     author='Matthias Fey',
     author_email='matthias.fey@tu-dortmund.de',
     url='https://github.com/rusty1s/pytorch_sparse',
diff --git a/torch_sparse/__init__.py b/torch_sparse/__init__.py
@@ -3,16 +3,18 @@
 
 import torch
 
-__version__ = '0.6.8'
+__version__ = '0.6.9'
+
+suffix = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 for library in [
         '_version', '_convert', '_diag', '_spmm', '_spspmm', '_metis', '_rw',
         '_saint', '_sample', '_relabel'
 ]:
     torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
-        library, [osp.dirname(__file__)]).origin)
+        f'{library}_{suffix}', [osp.dirname(__file__)]).origin)
 
-if torch.cuda.is_available() and torch.version.cuda:  # pragma: no cover
+if torch.cuda.is_available():  # pragma: no cover
     cuda_version = torch.ops.torch_sparse.cuda_version()
 
     if cuda_version == -1:
