Add PyTorch 2.1.0 support (#256)

pyg-team · Oct 5, 2023 · dda98fa · dda98fa
1 parent 42b8f0c
commit dda98fa
Show file tree

Hide file tree

Showing 10 changed files with 64 additions and 41 deletions.
diff --git a/.github/actions/setup/action.yml b/.github/actions/setup/action.yml
@@ -6,7 +6,7 @@ inputs:
     default: '3.8'
   torch-version:
     required: false
-    default: '2.0.0'
+    default: '2.1.0'
   cuda-version:
     required: false
     default: cpu

diff --git a/.github/workflows/building.yml b/.github/workflows/building.yml
@@ -12,39 +12,37 @@ jobs:
       matrix:
         os: [ubuntu-20.04, macos-10.15]  # windows-2019
         python-version: ['3.8', '3.9', '3.10', '3.11']
-        torch-version: [1.11.0, 1.12.0, 1.13.0, 2.0.0]
-        cuda-version: ['cpu', 'cu113', 'cu115', 'cu116', 'cu117', 'cu118']
+        torch-version: [1.12.0, 1.13.0, 2.0.0, 2.1.0]
+        cuda-version: ['cpu', 'cu113', 'cu116', 'cu117', 'cu118', 'cu121']
         exclude:
-          - torch-version: 1.11.0
-            python-version: '3.11'
-          - torch-version: 1.11.0
-            cuda-version: 'cu116'
-          - torch-version: 1.11.0
-            cuda-version: 'cu117'
-          - torch-version: 1.11.0
-            cuda-version: 'cu118'
           - torch-version: 1.12.0
             python-version: '3.11'
-          - torch-version: 1.12.0
-            cuda-version: 'cu115'
           - torch-version: 1.12.0
             cuda-version: 'cu117'
           - torch-version: 1.12.0
             cuda-version: 'cu118'
+          - torch-version: 1.12.0
+            cuda-version: 'cu121'
           - torch-version: 1.13.0
             python-version: '3.11'
           - torch-version: 1.13.0
             cuda-version: 'cu113'
-          - torch-version: 1.13.0
-            cuda-version: 'cu115'
           - torch-version: 1.13.0
             cuda-version: 'cu118'
+          - torch-version: 1.13.0
+            cuda-version: 'cu121'
           - torch-version: 2.0.0
             cuda-version: 'cu113'
           - torch-version: 2.0.0
-            cuda-version: 'cu115'
+            cuda-version: 'cu116'
+          - torch-version: 1.13.0
+            cuda-version: 'cu121'
+          - torch-version: 2.1.0
+            cuda-version: 'cu113'
           - torch-version: 2.0.0
             cuda-version: 'cu116'
+          - torch-version: 1.13.0
+            cuda-version: 'cu117'
           - os: macos-10.15
             cuda-version: 'cu113'
           - os: macos-10.15

diff --git a/.github/workflows/cuda/Linux-env.sh b/.github/workflows/cuda/Linux-env.sh
@@ -1,29 +1,40 @@
 #!/bin/bash
 
 case ${1} in
+  cu121)
+    export FORCE_CUDA=1
+    export PATH=/usr/local/cuda-12.1/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0"
+    ;;
   cu118)
     export FORCE_CUDA=1
     export PATH=/usr/local/cuda-11.8/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0"
     ;;
   cu117)
     export FORCE_CUDA=1
     export PATH=/usr/local/cuda-11.7/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
     ;;
   cu116)
     export FORCE_CUDA=1
     export PATH=/usr/local/cuda-11.6/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
     ;;
   cu115)
     export FORCE_CUDA=1
     export PATH=/usr/local/cuda-11.5/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
     ;;
   cu113)
     export FORCE_CUDA=1
     export PATH=/usr/local/cuda-11.3/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
     ;;
   cu102)
     export FORCE_CUDA=1
     export PATH=/usr/local/cuda-10.2/bin:${PATH}
+    export TORCH_CUDA_ARCH_LIST="3.5;5.0+PTX;6.0;7.0;7.5"
     ;;
   *)
     ;;

diff --git a/.github/workflows/cuda/Linux.sh b/.github/workflows/cuda/Linux.sh
@@ -3,6 +3,12 @@
 OS=ubuntu2004
 
 case ${1} in
+  cu121)
+    CUDA=12.1
+    APT_KEY=${OS}-${CUDA/./-}-local
+    FILENAME=cuda-repo-${APT_KEY}_${CUDA}.1-530.30.02-1_amd64.deb
+    URL=https://developer.download.nvidia.com/compute/cuda/${CUDA}.1/local_installers
+    ;;
   cu118)
     CUDA=11.8
     APT_KEY=${OS}-${CUDA/./-}-local
@@ -50,7 +56,7 @@ sudo mv cuda-${OS}.pin /etc/apt/preferences.d/cuda-repository-pin-600
 wget -nv ${URL}/${FILENAME}
 sudo dpkg -i ${FILENAME}
 
-if [ "${1}" = "cu117" ] || [ "${1}" = "cu118" ]; then
+if [ "${1}" = "cu117" ] || [ "${1}" = "cu118" ] || [ "${1}" = "cu121" ]; then
   sudo cp /var/cuda-repo-${APT_KEY}/cuda-*-keyring.gpg /usr/share/keyrings/
 else
   sudo apt-key add /var/cuda-repo-${APT_KEY}/7fa2af80.pub

diff --git a/.github/workflows/cuda/Windows-env.sh b/.github/workflows/cuda/Windows-env.sh
@@ -1,6 +1,10 @@
 #!/bin/bash
 
 case ${1} in
+  cu121)
+    export FORCE_CUDA=1
+    export PATH=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v12.1/bin:${PATH}
+    ;;
   cu118)
     export FORCE_CUDA=1
     export PATH=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v11.8/bin:${PATH}

diff --git a/.github/workflows/cuda/Windows.sh b/.github/workflows/cuda/Windows.sh
@@ -6,6 +6,11 @@ curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq
 7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32"
 
 case ${1} in
+  cu121)
+    CUDA_SHORT=11.8
+    CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.0/local_installers
+    CUDA_FILE=cuda_${CUDA_SHORT}.1_531.14_windows.exe
+    ;;
   cu118)
     CUDA_SHORT=11.8
     CUDA_URL=https://developer.download.nvidia.com/compute/cuda/${CUDA_SHORT}.0/local_installers

diff --git a/.github/workflows/install.yml b/.github/workflows/install.yml
@@ -13,7 +13,7 @@ jobs:
 
     strategy:
       matrix:
-        cuda-version: ['cpu', 'cu117']
+        cuda-version: ['cpu', 'cu121']
 
     steps:
       - name: Checkout repository

diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml
@@ -16,39 +16,37 @@ jobs:
       matrix:
         os: [ubuntu-20.04]
         python-version: ['3.8', '3.9', '3.10', '3.11']
-        torch-version: [1.11.0, 1.12.0, 1.13.0, 2.0.0]
-        cuda-version: ['cpu', 'cu113', 'cu115', 'cu116', 'cu117', 'cu118']
+        torch-version: [1.12.0, 1.13.0, 2.0.0, 2.1.0]
+        cuda-version: ['cpu', 'cu113', 'cu116', 'cu117', 'cu118', 'cu121']
         exclude:
-          - torch-version: 1.11.0
-            python-version: '3.11'
-          - torch-version: 1.11.0
-            cuda-version: 'cu116'
-          - torch-version: 1.11.0
-            cuda-version: 'cu117'
-          - torch-version: 1.11.0
-            cuda-version: 'cu118'
           - torch-version: 1.12.0
             python-version: '3.11'
-          - torch-version: 1.12.0
-            cuda-version: 'cu115'
           - torch-version: 1.12.0
             cuda-version: 'cu117'
           - torch-version: 1.12.0
             cuda-version: 'cu118'
+          - torch-version: 1.12.0
+            cuda-version: 'cu121'
           - torch-version: 1.13.0
             python-version: '3.11'
           - torch-version: 1.13.0
             cuda-version: 'cu113'
-          - torch-version: 1.13.0
-            cuda-version: 'cu115'
           - torch-version: 1.13.0
             cuda-version: 'cu118'
+          - torch-version: 1.13.0
+            cuda-version: 'cu121'
           - torch-version: 2.0.0
             cuda-version: 'cu113'
           - torch-version: 2.0.0
-            cuda-version: 'cu115'
+            cuda-version: 'cu116'
+          - torch-version: 1.13.0
+            cuda-version: 'cu121'
+          - torch-version: 2.1.0
+            cuda-version: 'cu113'
           - torch-version: 2.0.0
             cuda-version: 'cu116'
+          - torch-version: 1.13.0
+            cuda-version: 'cu117'
 
     steps:
       - name: Checkout repository

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [0.3.0] - 2023-MM-DD
 ### Added
+- Added PyTorch 2.1 support ([#256](https://github.com/pyg-team/pyg-lib/pull/256))
 - Added low-level support for distributed neighborhood sampling ([#246](https://github.com/pyg-team/pyg-lib/pull/246), [#252](https://github.com/pyg-team/pyg-lib/pull/252), [#253](https://github.com/pyg-team/pyg-lib/pull/253), [#254](https://github.com/pyg-team/pyg-lib/pull/254))
 - Added support for homogeneous and heterogeneous biased neighborhood sampling ([#247](https://github.com/pyg-team/pyg-lib/pull/247), [#251](https://github.com/pyg-team/pyg-lib/pull/251))
 - Added dispatch for XPU device in `index_sort` ([#243](https://github.com/pyg-team/pyg-lib/pull/243))

diff --git a/README.md b/README.md
@@ -26,31 +26,31 @@ pip install pyg-lib -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html
 
 where
 
-* `${TORCH}` should be replaced by either `1.11.0`, `1.12.0`, `1.13.0` or `2.0.0`
-* `${CUDA}` should be replaced by either `cpu`, `cu102`, `cu113`, `cu115`, `cu116`, `cu117` or `cu118`
+* `${TORCH}` should be replaced by either `1.12.0`, `1.13.0`, `2.0.0` or `2.1.0`
+* `${CUDA}` should be replaced by either `cpu`, `cu102`, `cu113`, `cu116`, `cu117`, `cu118` or `cu121`
 
 The following combinations are supported:
 
-| PyTorch 2.0  | `cpu` | `cu102` | `cu113` | `cu115` | `cu116` | `cu117` | `cu118` |
+| PyTorch 2.1  | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` |
 |--------------|-------|---------|---------|---------|---------|---------|---------|
 | **Linux**    | ✅    |         |         |         |         | ✅      | ✅      |
 | **Windows**  |       |         |         |         |         |         |         |
 | **macOS**    | ✅    |         |         |         |         |         |         |
 
-| PyTorch 1.13 | `cpu` | `cu102` | `cu113` | `cu115` | `cu116` | `cu117` | `cu118` |
+| PyTorch 2.0  | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` |
 |--------------|-------|---------|---------|---------|---------|---------|---------|
 | **Linux**    | ✅    |         |         |         | ✅      | ✅      |         |
 | **Windows**  |       |         |         |         |         |         |         |
 | **macOS**    | ✅    |         |         |         |         |         |         |
 
-| PyTorch 1.12 | `cpu` | `cu102` | `cu113` | `cu115` | `cu116` | `cu117` | `cu118` |
+| PyTorch 1.13 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` |
 |--------------|-------|---------|---------|---------|---------|---------|---------|
-| **Linux**    | ✅    | ✅      | ✅      |         | ✅      |         |         |
+| **Linux**    | ✅    |         |         | ✅      | ✅      |         |         |
 | **Windows**  |       |         |         |         |         |         |         |
 | **macOS**    | ✅    |         |         |         |         |         |         |
 
-| PyTorch 1.11 | `cpu` | `cu102` | `cu113` | `cu115` | `cu116` | `cu117` | `cu118` |
-|--------------|-------|---------|---------|---------|---------|---------|---------|
+| PyTorch 1.12 | `cpu` | `cu102` | `cu113` | `cu116` | `cu117` | `cu118` | `cu121` |
+|--------------|-------|---------|---------|---------|---------|---------| --------|
 | **Linux**    | ✅    | ✅      | ✅      | ✅      |         |         |         |
 | **Windows**  |       |         |         |         |         |         |         |
 | **macOS**    | ✅    |         |         |         |         |         |         |