From 50e2af6f4ab6178d8e6bcc531e7c26835e919570 Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Sun, 22 Oct 2023 16:09:34 +0800
Subject: [PATCH] support torch 2.1

---
 .github/workflows/build-doc.yml               |  2 +-
 .github/workflows/macos-cpu-wheels.yml        |  4 ++--
 .github/workflows/publish_to_pypi.yml         |  2 +-
 .github/workflows/run-tests-macos-cpu.yml     |  6 ++---
 .github/workflows/run-tests-ubuntu-cpu.yml    | 10 ++-------
 .github/workflows/run-tests-ubuntu-cuda.yml   | 14 +++---------
 .github/workflows/run-tests-windows-cpu.yml   |  6 ++---
 .github/workflows/run-tests-windows-cuda.yml  |  6 ++---
 .github/workflows/style_check.yml             |  6 ++---
 .github/workflows/test-wheels.yml             |  2 +-
 .github/workflows/ubuntu-cpu-wheels.yml       |  4 ++--
 .github/workflows/windows-x64-cpu-wheels.yml  |  4 ++--
 .../github_actions/generate_build_matrix.py   |  9 +++++++-
 scripts/github_actions/install_cuda.sh        | 19 +++++++++++++++-
 scripts/github_actions/install_cudnn.sh       |  6 +++++
 scripts/github_actions/install_torch.sh       | 22 +++++++++++++++----
 16 files changed, 76 insertions(+), 46 deletions(-)

diff --git a/.github/workflows/build-doc.yml b/.github/workflows/build-doc.yml
index e10eb8b..43e76c4 100644
--- a/.github/workflows/build-doc.yml
+++ b/.github/workflows/build-doc.yml
@@ -35,7 +35,7 @@ jobs:
         python-version: [3.8]
     steps:
       # refer to https://github.com/actions/checkout
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/macos-cpu-wheels.yml b/.github/workflows/macos-cpu-wheels.yml
index bbb8ec6..ec8fc03 100644
--- a/.github/workflows/macos-cpu-wheels.yml
+++ b/.github/workflows/macos-cpu-wheels.yml
@@ -17,7 +17,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -38,7 +38,7 @@ jobs:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/publish_to_pypi.yml b/.github/workflows/publish_to_pypi.yml
index 5c93a29..3b48f39 100644
--- a/.github/workflows/publish_to_pypi.yml
+++ b/.github/workflows/publish_to_pypi.yml
@@ -26,7 +26,7 @@ jobs:
   pypi:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/run-tests-macos-cpu.yml b/.github/workflows/run-tests-macos-cpu.yml
index 6707c7e..a61129d 100644
--- a/.github/workflows/run-tests-macos-cpu.yml
+++ b/.github/workflows/run-tests-macos-cpu.yml
@@ -32,7 +32,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -45,14 +45,14 @@ jobs:
 
   run_tests_macos_cpu:
     needs: generate_build_matrix
-    runs-on: macos-10.15
+    runs-on: macos-latest
     strategy:
       fail-fast: false
       matrix:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/run-tests-ubuntu-cpu.yml b/.github/workflows/run-tests-ubuntu-cpu.yml
index 4337cce..0e4cb73 100644
--- a/.github/workflows/run-tests-ubuntu-cpu.yml
+++ b/.github/workflows/run-tests-ubuntu-cpu.yml
@@ -32,7 +32,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -52,7 +52,7 @@ jobs:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
@@ -61,12 +61,6 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Install GCC 7
-        run: |
-          sudo apt-get install -y gcc-7 g++-7
-          echo "CC=/usr/bin/gcc-7" >> $GITHUB_ENV
-          echo "CXX=/usr/bin/g++-7" >> $GITHUB_ENV
-
       - name: Install PyTorch ${{ matrix.torch }}
         shell: bash
         run: |
diff --git a/.github/workflows/run-tests-ubuntu-cuda.yml b/.github/workflows/run-tests-ubuntu-cuda.yml
index 51251b5..a168ec8 100644
--- a/.github/workflows/run-tests-ubuntu-cuda.yml
+++ b/.github/workflows/run-tests-ubuntu-cuda.yml
@@ -32,7 +32,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -45,14 +45,14 @@ jobs:
 
   run_tests_ubuntu_cuda:
     needs: generate_build_matrix
-    runs-on: ubuntu-18.04
+    runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
@@ -61,14 +61,6 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Install GCC 7
-        if: startsWith(matrix.os, 'ubuntu')
-        run: |
-          sudo apt-get install -y gcc-7 g++-7
-          echo "CC=/usr/bin/gcc-7" >> $GITHUB_ENV
-          echo "CXX=/usr/bin/g++-7" >> $GITHUB_ENV
-
-
       - name: Install CUDA Toolkit ${{ matrix.cuda }}
         shell: bash
         env:
diff --git a/.github/workflows/run-tests-windows-cpu.yml b/.github/workflows/run-tests-windows-cpu.yml
index f735ac3..4525bad 100644
--- a/.github/workflows/run-tests-windows-cpu.yml
+++ b/.github/workflows/run-tests-windows-cpu.yml
@@ -32,7 +32,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -46,14 +46,14 @@ jobs:
   run_tests_windows_cpu:
     # see https://github.com/actions/virtual-environments/blob/win19/20210525.0/images/win/Windows2019-Readme.md
     needs: generate_build_matrix
-    runs-on: windows-2019
+    runs-on: windows-latest
     strategy:
       fail-fast: false
       matrix:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/run-tests-windows-cuda.yml b/.github/workflows/run-tests-windows-cuda.yml
index 6d04dd5..9220e12 100644
--- a/.github/workflows/run-tests-windows-cuda.yml
+++ b/.github/workflows/run-tests-windows-cuda.yml
@@ -33,7 +33,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -46,14 +46,14 @@ jobs:
 
   run_tests_windows_cuda:
     needs: generate_build_matrix
-    runs-on: windows-2019
+    runs-on: windows-latest
     strategy:
       fail-fast: false
       matrix:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/style_check.yml b/.github/workflows/style_check.yml
index 462a2db..4ee8aec 100644
--- a/.github/workflows/style_check.yml
+++ b/.github/workflows/style_check.yml
@@ -29,12 +29,12 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os: [ubuntu-18.04, macos-10.15]
-        python-version: [3.7, 3.8, 3.9]
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.8"]
       fail-fast: false
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/test-wheels.yml b/.github/workflows/test-wheels.yml
index 40a06eb..9d3ea70 100644
--- a/.github/workflows/test-wheels.yml
+++ b/.github/workflows/test-wheels.yml
@@ -20,7 +20,7 @@ jobs:
         python-version: ["3.8", "3.9", "3.10"]
     steps:
       # refer to https://github.com/actions/checkout
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/ubuntu-cpu-wheels.yml b/.github/workflows/ubuntu-cpu-wheels.yml
index 6f9507a..36246fa 100644
--- a/.github/workflows/ubuntu-cpu-wheels.yml
+++ b/.github/workflows/ubuntu-cpu-wheels.yml
@@ -17,7 +17,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -38,7 +38,7 @@ jobs:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/.github/workflows/windows-x64-cpu-wheels.yml b/.github/workflows/windows-x64-cpu-wheels.yml
index 17f1793..7f8f233 100644
--- a/.github/workflows/windows-x64-cpu-wheels.yml
+++ b/.github/workflows/windows-x64-cpu-wheels.yml
@@ -17,7 +17,7 @@ jobs:
     outputs:
       matrix: ${{ steps.set-matrix.outputs.matrix }}
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Generating build matrix
@@ -38,7 +38,7 @@ jobs:
         ${{ fromJson(needs.generate_build_matrix.outputs.matrix) }}
 
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
 
diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
index df1dc24..151a005 100755
--- a/scripts/github_actions/generate_build_matrix.py
+++ b/scripts/github_actions/generate_build_matrix.py
@@ -147,9 +147,16 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
             if not for_windows
             else ["11.7.1", "11.8.0"],
         },
+        "2.1.0": {
+            "python-version": ["3.8", "3.9", "3.10", "3.11"],
+            "cuda": ["11.8", "12.1"]  # default 12.1
+            if not for_windows
+            else ["11.8.0", "12.1.0"],
+        },
+        # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/windows-links.ts
     }
     if test_only_latest_torch:
-        latest = "2.0.1"
+        latest = "2.1.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
diff --git a/scripts/github_actions/install_cuda.sh b/scripts/github_actions/install_cuda.sh
index f7a669a..3d47f4c 100755
--- a/scripts/github_actions/install_cuda.sh
+++ b/scripts/github_actions/install_cuda.sh
@@ -49,6 +49,12 @@ case "$cuda" in
   11.7)
     url=https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.65.01_linux.run
     ;;
+  11.8)
+    url=https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
+    ;;
+  12.1)
+    url=https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run
+    ;;
   *)
     echo "Unknown cuda version: $cuda"
     exit 1
@@ -63,10 +69,21 @@ retry curl -LSs -O $url
 filename=$(basename $url)
 echo "filename: $filename"
 chmod +x ./$filename
-sudo ./$filename --toolkit --silent
+
+ls -lh
+ls -lh /usr/local
+
+sudo ./$filename \
+  --silent \
+  --toolkit \
+  --no-opengl-libs \
+  --no-drm \
+  --no-man-page
+
 rm -fv ./$filename
 
 export CUDA_HOME=/usr/local/cuda
 export PATH=$CUDA_HOME/bin:$PATH
 export LD_LIBRARY_PATH=$CUDA_HOME/lib:$LD_LIBRARY_PATH
 export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH
+ls -lh $CUDA_HOME
diff --git a/scripts/github_actions/install_cudnn.sh b/scripts/github_actions/install_cudnn.sh
index d57018c..527e9a7 100755
--- a/scripts/github_actions/install_cudnn.sh
+++ b/scripts/github_actions/install_cudnn.sh
@@ -42,6 +42,12 @@ case $cuda in
   11.7)
     filename=cudnn-11.3-linux-x64-v8.2.0.53.tgz
     ;;
+  11.8)
+    filename=cudnn-11.3-linux-x64-v8.2.0.53.tgz
+    ;;
+  12.1)
+    filename=cudnn-linux-x86_64-8.9.5.29_cuda12-archive.tar.xz
+    ;;
   *)
     echo "Unsupported cuda version: $cuda"
     exit 1
diff --git a/scripts/github_actions/install_torch.sh b/scripts/github_actions/install_torch.sh
index c63138f..1209ed6 100755
--- a/scripts/github_actions/install_torch.sh
+++ b/scripts/github_actions/install_torch.sh
@@ -13,11 +13,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-set -x
 
-torch=$TORCH_VERSION
-cuda=$CUDA_VERSION
-echo "torch: $torch, cuda: $cuda"
+if [ x"$TORCH_VERSION" != x"" ] && [ x"$CUDA_VERSION" != x"" ]; then
+    torch=$TORCH_VERSION
+    cuda=$CUDA_VERSION
+fi
+
 case ${torch} in
   1.5.*)
     case ${cuda} in
@@ -171,6 +172,19 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.1.*)
+    case ${cuda} in
+      11.8)
+        package="torch==${torch}+cu118"
+        url=https://download.pytorch.org/whl/torch_stable.html
+        ;;
+      12.1)
+        package="torch==${torch}"
+        # Leave it empty to use PyPI.
+        url=
+        ;;
+    esac
+    ;;
   *)
     echo "Unsupported PyTorch version: ${torch}"
     exit 1