From cb8d7f31efdd86db840827e21c922e7084c99b93 Mon Sep 17 00:00:00 2001
From: atalman <atalman@fb.com>
Date: Mon, 21 Mar 2022 15:16:25 -0700
Subject: [PATCH] Adding windows cuda 11.6 changes

---
 .github/workflows/build-magma-windows.yml |  2 +-
 conda/pytorch-nightly/bld.bat             |  4 ++
 windows/cuda116.bat                       | 58 +++++++++++++++++++++++
 windows/internal/build_magma.bat          | 13 +----
 windows/internal/cuda_install.bat         | 28 +++++++++++
 windows/internal/smoke_test.bat           |  4 +-
 6 files changed, 95 insertions(+), 14 deletions(-)
 create mode 100644 windows/cuda116.bat

diff --git a/.github/workflows/build-magma-windows.yml b/.github/workflows/build-magma-windows.yml
index 6f82a1b24..76e973d17 100644
--- a/.github/workflows/build-magma-windows.yml
+++ b/.github/workflows/build-magma-windows.yml
@@ -15,7 +15,7 @@ jobs:
     runs-on: windows-2019
     strategy:
       matrix:
-        cuda_version: ["115"]
+        cuda_version: ["116", "115"]
         config: ["Release", "Debug"]
     env:
       CUDA_VERSION: ${{ matrix.cuda_version }}
diff --git a/conda/pytorch-nightly/bld.bat b/conda/pytorch-nightly/bld.bat
index 7a4bd7739..db1671bf0 100644
--- a/conda/pytorch-nightly/bld.bat
+++ b/conda/pytorch-nightly/bld.bat
@@ -30,6 +30,10 @@ if "%desired_cuda%" == "11.5" (
     set TORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST%;6.0;6.1;7.0;7.5;8.0;8.6
     set TORCH_NVCC_FLAGS=-Xfatbin -compress-all --threads 2
 )
+if "%desired_cuda%" == "11.6" (
+    set TORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST%;6.0;6.1;7.0;7.5;8.0;8.6
+    set TORCH_NVCC_FLAGS=-Xfatbin -compress-all --threads 2
+)
 
 :cuda_flags_end
 
diff --git a/windows/cuda116.bat b/windows/cuda116.bat
new file mode 100644
index 000000000..a5c73039f
--- /dev/null
+++ b/windows/cuda116.bat
@@ -0,0 +1,58 @@
+@echo off
+
+set MODULE_NAME=pytorch
+
+IF NOT EXIST "setup.py" IF NOT EXIST "%MODULE_NAME%" (
+    call internal\clone.bat
+    cd ..
+) ELSE (
+    call internal\clean.bat
+)
+IF ERRORLEVEL 1 goto :eof
+
+call internal\check_deps.bat
+IF ERRORLEVEL 1 goto :eof
+
+REM Check for optional components
+
+set USE_CUDA=
+set CMAKE_GENERATOR=Visual Studio 15 2017 Win64
+
+IF "%NVTOOLSEXT_PATH%"=="" (
+    IF EXIST "C:\Program Files\NVIDIA Corporation\NvToolsExt\lib\x64\nvToolsExt64_1.lib"  (
+        set NVTOOLSEXT_PATH=C:\Program Files\NVIDIA Corporation\NvToolsExt
+    ) ELSE (
+        echo NVTX ^(Visual Studio Extension ^for CUDA^) ^not installed, failing
+        exit /b 1
+    )
+)
+
+IF "%CUDA_PATH_V116%"=="" (
+    IF EXIST "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6\bin\nvcc.exe" (
+        set "CUDA_PATH_V115=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.6"
+    ) ELSE (
+        echo CUDA 11.6 not found, failing
+        exit /b 1
+    )
+)
+
+IF "%BUILD_VISION%" == "" (
+    set TORCH_CUDA_ARCH_LIST=3.7+PTX;5.0;6.0;6.1;7.0;7.5;8.0;8.6
+    set TORCH_NVCC_FLAGS=-Xfatbin -compress-all
+) ELSE (
+    set NVCC_FLAGS=-D__CUDA_NO_HALF_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_35,code=sm_35 -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_86,code=compute_86
+)
+
+set "CUDA_PATH=%CUDA_PATH_V116%"
+set "PATH=%CUDA_PATH_V116%\bin;%PATH%"
+
+:optcheck
+
+call internal\check_opts.bat
+IF ERRORLEVEL 1 goto :eof
+
+call internal\copy.bat
+IF ERRORLEVEL 1 goto :eof
+
+call internal\setup.bat
+IF ERRORLEVEL 1 goto :eof
diff --git a/windows/internal/build_magma.bat b/windows/internal/build_magma.bat
index 795acb27c..b8a2f3377 100644
--- a/windows/internal/build_magma.bat
+++ b/windows/internal/build_magma.bat
@@ -36,18 +36,7 @@ mkdir build && cd build
 
 set GPU_TARGET=All
 set CUDA_ARCH_LIST= -gencode arch=compute_37,code=sm_37 -gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70
-
-IF "%CUVER_NODOT%" == "110" (
-  set "CUDA_ARCH_LIST=%CUDA_ARCH_LIST% -gencode arch=compute_80,code=sm_80"
-)
-
-IF "%CUVER_NODOT%" == "113" (
-  set "CUDA_ARCH_LIST=%CUDA_ARCH_LIST% -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86"
-)
-
-IF "%CUVER_NODOT%" == "115" (
-  set "CUDA_ARCH_LIST=%CUDA_ARCH_LIST% -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86"
-)
+set "CUDA_ARCH_LIST=%CUDA_ARCH_LIST% -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86"
 
 set CC=cl.exe
 set CXX=cl.exe
diff --git a/windows/internal/cuda_install.bat b/windows/internal/cuda_install.bat
index 58a0a9802..ce8a8ccf4 100644
--- a/windows/internal/cuda_install.bat
+++ b/windows/internal/cuda_install.bat
@@ -22,6 +22,7 @@ if exist "C:\\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR
 if %CUDA_VER% EQU 102 goto cuda102
 if %CUDA_VER% EQU 113 goto cuda113
 if %CUDA_VER% EQU 115 goto cuda115
+if %CUDA_VER% EQU 116 goto cuda116
 
 echo CUDA %CUDA_VERSION_STR% is not supported
 exit /b 1
@@ -89,6 +90,33 @@ xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
 
 goto cuda_common
 
+:cuda116
+
+set CUDA_INSTALL_EXE=cuda_11.6.0_511.23_windows.exe
+if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
+    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    if errorlevel 1 exit /b 1
+    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    set "ARGS=thrust_11.6 nvcc_11.6 cuobjdump_11.6 nvprune_11.6 nvprof_11.6 cupti_11.6 cublas_11.6 cublas_dev_11.6 cudart_11.6 cufft_11.6 cufft_dev_11.6 curand_11.6 curand_dev_11.6 cusolver_11.6 cusolver_dev_11.6 cusparse_11.6 cusparse_dev_11.6 npp_11.6 npp_dev_11.6 nvrtc_11.6 nvrtc_dev_11.6 nvml_dev_11.6"
+)
+
+set CUDNN_FOLDER=cudnn-windows-x86_64-8.3.2.44_cuda11.5-archive
+set CUDNN_LIB_FOLDER="lib"
+set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
+if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
+    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+    if errorlevel 1 exit /b 1
+    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+)
+
+@REM Cuda 8.3+ required zlib to be installed on the path
+echo Installing ZLIB dlls
+curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
+7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
+xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
+
+goto cuda_common
+
 :cuda_common
 :: NOTE: We only install CUDA if we don't have it installed already.
 :: With GHA runners these should be pre-installed as part of our AMI process
diff --git a/windows/internal/smoke_test.bat b/windows/internal/smoke_test.bat
index fc142a1d2..8e626cf96 100644
--- a/windows/internal/smoke_test.bat
+++ b/windows/internal/smoke_test.bat
@@ -72,7 +72,9 @@ set "CONDA_EXTRA_ARGS="
 if "%CUDA_VERSION%" == "115" (
     set "CONDA_EXTRA_ARGS=-c=nvidia"
 )
-
+if "%CUDA_VERSION%" == "116" (
+    set "CONDA_EXTRA_ARGS=-c=nvidia"
+)
 
 rmdir /s /q conda
 del miniconda.exe