Merge pull request #162 from RomeoV/master

Let torch determine correct cuda architecture

Author: rusty1s

.github/workflows/building.yml

@@ -70,7 +70,7 @@ jobs:
         if: ${{ matrix.cuda-version != 'cpu' }}
         run: |
           source .github/workflows/cuda/${{ matrix.cuda-version }}-${{ runner.os }}-env.sh
-          FORCE_CUDA=1 WITH_METIS=1 pip install -e .
+          FORCE_CUDA=1 WITH_METIS=1 TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX" pip install -e .
         shell:
           bash
 

CMakeLists.txt

@@ -9,7 +9,7 @@ if(WITH_CUDA)
   enable_language(CUDA)
   add_definitions(-D__CUDA_NO_HALF_OPERATORS__)
   add_definitions(-DWITH_CUDA)
-  set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -arch=sm_35 --expt-relaxed-constexpr")
+  set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
 endif()
 
 find_package(Python3 COMPONENTS Development)

conda/pytorch-sparse/meta.yaml

@@ -27,6 +27,7 @@ build:
   script_env:
     - FORCE_CUDA
     - WITH_METIS=1
+    - TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX"
   preserve_egg_dir: True
 
 test:

setup.py

@@ -62,7 +62,7 @@ def get_extensions():
             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']
+            nvcc_flags += ['--expt-relaxed-constexpr', '-O2']
             extra_compile_args['nvcc'] = nvcc_flags
 
             if sys.platform == 'win32':