From b9e1e81383e3a03bd5c934ed0f4ace8a1739c1ac Mon Sep 17 00:00:00 2001
From: pbialecki <piotr.bialecki@hotmail.de>
Date: Thu, 30 Mar 2023 02:49:28 -0700
Subject: [PATCH] build libtorch and manywheel for 12.1

---
 .github/workflows/build-libtorch-images.yml  | 2 +-
 .github/workflows/build-manywheel-images.yml | 2 +-
 libtorch/Dockerfile                          | 4 ++++
 libtorch/build_all_docker.sh                 | 2 +-
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/build-libtorch-images.yml b/.github/workflows/build-libtorch-images.yml
index e2493679a..d95f9208c 100644
--- a/.github/workflows/build-libtorch-images.yml
+++ b/.github/workflows/build-libtorch-images.yml
@@ -28,7 +28,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        cuda_version: ["11.8", "11.7"]
+        cuda_version: ["12.1", "11.8", "11.7"]
     env:
       GPU_ARCH_TYPE: cuda
       GPU_ARCH_VERSION: ${{ matrix.cuda_version }}
diff --git a/.github/workflows/build-manywheel-images.yml b/.github/workflows/build-manywheel-images.yml
index 18511f027..b171edb80 100644
--- a/.github/workflows/build-manywheel-images.yml
+++ b/.github/workflows/build-manywheel-images.yml
@@ -30,7 +30,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        cuda_version: ["11.8", "11.7"]
+        cuda_version: ["12.1", "11.8", "11.7"]
     env:
       GPU_ARCH_TYPE: cuda
       GPU_ARCH_VERSION: ${{ matrix.cuda_version }}
diff --git a/libtorch/Dockerfile b/libtorch/Dockerfile
index 6a461b808..cfe1208bd 100644
--- a/libtorch/Dockerfile
+++ b/libtorch/Dockerfile
@@ -52,6 +52,10 @@ FROM cuda as cuda11.8
 RUN bash ./install_cuda.sh 11.8
 RUN bash ./install_magma.sh 11.8
 
+FROM cuda as cuda12.1
+RUN bash ./install_cuda.sh 12.1
+RUN bash ./install_magma.sh 12.1
+
 FROM cpu as rocm
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
diff --git a/libtorch/build_all_docker.sh b/libtorch/build_all_docker.sh
index bebb1728a..3c8aff9cb 100755
--- a/libtorch/build_all_docker.sh
+++ b/libtorch/build_all_docker.sh
@@ -4,7 +4,7 @@ set -eou pipefail
 
 TOPDIR=$(git rev-parse --show-toplevel)
 
-for cuda_version in 11.8 11.7; do
+for cuda_version in 12.1 11.8 11.7; do
     GPU_ARCH_TYPE=cuda GPU_ARCH_VERSION="${cuda_version}" "${TOPDIR}/libtorch/build_docker.sh"
 done