From a97794e013e0f02c9e9e3d48623ff613ace2c85d Mon Sep 17 00:00:00 2001
From: pbialecki <pbialecki@nvidia.com>
Date: Fri, 25 Mar 2022 13:48:09 -0700
Subject: [PATCH 1/3] Build libtorch and manywheel for 11.6

---
 .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 950d1fd1c..1045bd277 100644
--- a/.github/workflows/build-libtorch-images.yml
+++ b/.github/workflows/build-libtorch-images.yml
@@ -28,7 +28,7 @@ jobs:
     runs-on: ubuntu-18.04
     strategy:
       matrix:
-        cuda_version: ["11.5", "11.3", "10.2"]
+        cuda_version: ["11.6", "11.5", "11.3", "10.2"]
     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 13882ea1f..ec755d0c2 100644
--- a/.github/workflows/build-manywheel-images.yml
+++ b/.github/workflows/build-manywheel-images.yml
@@ -28,7 +28,7 @@ jobs:
     runs-on: ubuntu-18.04
     strategy:
       matrix:
-        cuda_version: ["11.5", "11.3", "10.2"]
+        cuda_version: ["11.6", "11.5", "11.3", "10.2"]
     env:
       GPU_ARCH_TYPE: cuda
       GPU_ARCH_VERSION: ${{ matrix.cuda_version }}
diff --git a/libtorch/Dockerfile b/libtorch/Dockerfile
index 86eef5602..1c4719171 100644
--- a/libtorch/Dockerfile
+++ b/libtorch/Dockerfile
@@ -56,6 +56,10 @@ FROM cuda as cuda11.5
 RUN bash ./install_cuda.sh 11.5
 RUN bash ./install_magma.sh 11.5
 
+FROM cuda as cuda11.6
+RUN bash ./install_cuda.sh 11.6
+RUN bash ./install_magma 11.6
+
 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 57f26bc8c..9d85fba84 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.5 11.3 10.2; do
+for cuda_version in 11.6 11.5 11.3 10.2; do
     GPU_ARCH_TYPE=cuda GPU_ARCH_VERSION="${cuda_version}" "${TOPDIR}/libtorch/build_docker.sh"
 done
 

From 910c2aafa68aab8a5dab2657eaeaaa6e24b98a2e Mon Sep 17 00:00:00 2001
From: Andrey Talman <atalman@fb.com>
Date: Mon, 28 Mar 2022 14:37:08 -0400
Subject: [PATCH 2/3] Make sure correct version of magma is installed

Magma version was updated in conda. Hence this update.
---
 common/install_magma.sh | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/common/install_magma.sh b/common/install_magma.sh
index 0e5212f0b..ea2548140 100644
--- a/common/install_magma.sh
+++ b/common/install_magma.sh
@@ -7,7 +7,14 @@ MAGMA_VERSION="2.5.2"
 function do_install() {
     cuda_version=$1
     cuda_version_nodot=${1/./}
-    magma_archive="magma-cuda${cuda_version_nodot}-${MAGMA_VERSION}-1.tar.bz2"
+    
+    if [[ ${cuda_version_nodot} == 116 ]]; then
+        MAGMA_VERSION="2.6.1"
+        magma_archive="magma-cuda${cuda_version_nodot}-${MAGMA_VERSION}-0.tar.bz2"
+    else
+        magma_archive="magma-cuda${cuda_version_nodot}-${MAGMA_VERSION}-1.tar.bz2"
+    fi
+    
     cuda_dir="/usr/local/cuda-${cuda_version}"
     (
         set -x

From da69191ab4037e8ad34936472a334fa812e30670 Mon Sep 17 00:00:00 2001
From: Andrey Talman <atalman@fb.com>
Date: Mon, 28 Mar 2022 14:52:26 -0400
Subject: [PATCH 3/3] Fix typo

Fix typo
---
 libtorch/Dockerfile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/libtorch/Dockerfile b/libtorch/Dockerfile
index 1c4719171..546a3181b 100644
--- a/libtorch/Dockerfile
+++ b/libtorch/Dockerfile
@@ -58,7 +58,7 @@ RUN bash ./install_magma.sh 11.5
 
 FROM cuda as cuda11.6
 RUN bash ./install_cuda.sh 11.6
-RUN bash ./install_magma 11.6
+RUN bash ./install_magma.sh 11.6
 
 FROM cpu as rocm
 ARG PYTORCH_ROCM_ARCH