From 8cb8035e87d758a466f0b02b2b9fb90a5a23149d Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 4 Oct 2023 22:52:25 +0200
Subject: [PATCH] python310Packages.torch: 2.0.1 -> 2.1.1

---
 .../python-modules/torch/default.nix          | 20 +++---
 .../torch/fix-cmake-cuda-toolkit.patch        | 65 +++++++++++++++++++
 2 files changed, 77 insertions(+), 8 deletions(-)
 create mode 100644 pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch

diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
index dae916a18525..49ed033bd6a1 100644
--- a/pkgs/development/python-modules/torch/default.nix
+++ b/pkgs/development/python-modules/torch/default.nix
@@ -20,6 +20,7 @@
   Accelerate, CoreServices, libobjc,
 
   # Propagated build inputs
+  fsspec,
   filelock,
   jinja2,
   networkx,
@@ -127,7 +128,7 @@ let
 in buildPythonPackage rec {
   pname = "torch";
   # Don't forget to update torch-bin to the same version.
-  version = "2.0.1";
+  version = "2.1.1";
   format = "setuptools";
 
   disabled = pythonOlder "3.8.0";
@@ -143,10 +144,13 @@ in buildPythonPackage rec {
     repo = "pytorch";
     rev = "refs/tags/v${version}";
     fetchSubmodules = true;
-    hash = "sha256-xUj77yKz3IQ3gd/G32pI4OhL3LoN1zS7eFg0/0nZp5I=";
+    hash = "sha256-01+uqHvPbQVXKLohGWfsCsZOjb7xmfjBKkTGUGMEdAI=";
   };
 
-  patches = lib.optionals (stdenv.isDarwin && stdenv.isx86_64) [
+  patches = lib.optionals cudaSupport [
+    ./fix-cmake-cuda-toolkit.patch
+  ]
+  ++ lib.optionals (stdenv.isDarwin && stdenv.isx86_64) [
     # pthreadpool added support for Grand Central Dispatch in April
     # 2020. However, this relies on functionality (DISPATCH_APPLY_AUTO)
     # that is available starting with macOS 10.13. However, our current
@@ -188,12 +192,11 @@ in buildPythonPackage rec {
         'message(FATAL_ERROR "Found NCCL header version and library version' \
         'message(WARNING "Found NCCL header version and library version'
   ''
-  # TODO(@connorbaker): Remove this patch after 2.1.0 lands.
+  # Remove PyTorch's FindCUDAToolkit.cmake and to use CMake's default.
+  # We do not remove the entirety of cmake/Modules_CUDA_fix because we need FindCUDNN.cmake.
   + lib.optionalString cudaSupport ''
-    substituteInPlace torch/utils/cpp_extension.py \
-      --replace \
-        "'8.6', '8.9'" \
-        "'8.6', '8.7', '8.9'"
+    rm cmake/Modules/FindCUDAToolkit.cmake
+    rm -rf cmake/Modules_CUDA_fix/{upstream,FindCUDA.cmake}
   ''
   # error: no member named 'aligned_alloc' in the global namespace; did you mean simply 'aligned_alloc'
   # This lib overrided aligned_alloc hence the error message. Tltr: his function is linkable but not in header.
@@ -377,6 +380,7 @@ in buildPythonPackage rec {
     pyyaml
 
     # From install_requires:
+    fsspec
     filelock
     typing-extensions
     sympy
diff --git a/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch b/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch
new file mode 100644
index 000000000000..508118cc2e8a
--- /dev/null
+++ b/pkgs/development/python-modules/torch/fix-cmake-cuda-toolkit.patch
@@ -0,0 +1,65 @@
+From aa7e7cfd5a5fe0c839406c1b197b81ef955a7ef5 Mon Sep 17 00:00:00 2001
+From: Gaetan Lepage <gaetan@glepage.com>
+Date: Thu, 23 Nov 2023 00:39:26 +0100
+Subject: [PATCH] fix-cmake-cuda-toolkit
+
+---
+ CMakeLists.txt               |  4 ----
+ cmake/public/cuda.cmake      | 12 +++++++++---
+ tools/setup_helpers/cmake.py |  2 ++
+ 3 files changed, 11 insertions(+), 7 deletions(-)
+
+diff --git a/CMakeLists.txt b/CMakeLists.txt
+index 3a48eaf4e29..3aaeef2593a 100644
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -1152,10 +1152,6 @@ if(BUILD_SHARED_LIBS)
+       ${PROJECT_SOURCE_DIR}/cmake/Modules_CUDA_fix
+       DESTINATION share/cmake/Caffe2/
+       COMPONENT dev)
+-  install(FILES
+-      ${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUDAToolkit.cmake
+-      DESTINATION share/cmake/Caffe2/
+-      COMPONENT dev)
+   install(FILES
+       ${PROJECT_SOURCE_DIR}/cmake/Modules/FindCUSPARSELT.cmake
+       DESTINATION share/cmake/Caffe2/
+diff --git a/cmake/public/cuda.cmake b/cmake/public/cuda.cmake
+index 32f3ba375b5..c8b666e4c78 100644
+--- a/cmake/public/cuda.cmake
++++ b/cmake/public/cuda.cmake
+@@ -62,9 +62,15 @@ cmake_policy(POP)
+ 
+ if(NOT CMAKE_CUDA_COMPILER_VERSION STREQUAL CUDAToolkit_VERSION OR
+     NOT CUDA_INCLUDE_DIRS STREQUAL CUDAToolkit_INCLUDE_DIR)
+-  message(FATAL_ERROR "Found two conflicting CUDA installs:\n"
+-                      "V${CMAKE_CUDA_COMPILER_VERSION} in '${CUDA_INCLUDE_DIRS}' and\n"
+-                      "V${CUDAToolkit_VERSION} in '${CUDAToolkit_INCLUDE_DIR}'")
++  if(CUDA_INCLUDE_DIRS IN_LIST CUDAToolkit_INCLUDE_DIR)
++    message(STATUS "CUDA_INCLUDE_DIRS is a substring of CUDAToolkit_INCLUDE_DIR. "
++                   "Setting CUDA_INCLUDE_DIRS to CUDAToolkit_INCLUDE_DIR.")
++    set(CUDA_INCLUDE_DIRS "${CUDAToolkit_INCLUDE_DIR}")
++  else()
++    message(FATAL_ERROR "Found two conflicting CUDA installs:\n"
++                        "V${CMAKE_CUDA_COMPILER_VERSION} in '${CUDA_INCLUDE_DIRS}' and\n"
++                        "V${CUDAToolkit_VERSION} in '${CUDAToolkit_INCLUDE_DIR}'")
++  endif()
+ endif()
+ 
+ if(NOT TARGET CUDA::nvToolsExt)
+diff --git a/tools/setup_helpers/cmake.py b/tools/setup_helpers/cmake.py
+index fb19b66dfba..3f83bef32fe 100644
+--- a/tools/setup_helpers/cmake.py
++++ b/tools/setup_helpers/cmake.py
+@@ -207,6 +207,8 @@ class CMake:
+                     "BUILDING_WITH_TORCH_LIBS",
+                     "CUDA_HOST_COMPILER",
+                     "CUDA_NVCC_EXECUTABLE",
++                    "CUDAToolkit_ROOT",
++                    "CUDAToolkit_INCLUDE_DIR",
+                     "CUDA_SEPARABLE_COMPILATION",
+                     "CUDNN_LIBRARY",
+                     "CUDNN_INCLUDE_DIR",
+-- 
+2.42.0
+