torch-2.9.0-rc2.diff

--- pkgs/python-modules/torch_2_8/default.nix   2025-09-16 11:12:40
+++ pkgs/python-modules/torch_2_9/default.nix   2025-09-16 11:12:47
@@ -4,6 +4,7 @@
   gcc11Stdenv,
   lib,
   fetchFromGitHub,
+  fetchpatch,
   buildPythonPackage,
   python,
   config,
@@ -133,8 +134,18 @@

   supportedTorchCudaCapabilities =
     let
-      # https://github.com/pytorch/pytorch/blob/release/2.8/.ci/manywheel/build_cuda.sh
+      # https://github.com/pytorch/pytorch/blob/release/2.9/.ci/manywheel/build_cuda.sh
       capsPerCudaVersion = {
+        "13.0" = [
+          "7.5"
+          "8.0"
+          "8.6"
+          "9.0"
+          "10.0"
+          "12.0"
+        ];
+        # NOTE: 12.9 does not seem to be in RC builds, check if needed for final release.
+        #       https://download.pytorch.org/whl/test/torch/
         "12.9" = [
           "7.0"
           "7.5"
@@ -190,6 +201,9 @@
     "gfx1030"
     "gfx1100"
     "gfx1101"
+    "gfx1102"
+    "gfx1200"
+    "gfx1201"
   ];

   # NOTE: The lists.subtractLists function is perhaps a bit unintuitive. It subtracts the elements
@@ -236,6 +250,7 @@
       hiprand
       hipsolver
       hipsparse
+      hipsparselt
       hsa-rocr
       miopen-hip
       rccl
@@ -269,8 +284,8 @@
     "Unsupported CUDA version" =
       cudaSupport
       && !(builtins.elem cudaPackages.cudaMajorVersion [
-        "11"
         "12"
+        "13"
       ]);
     "MPI cudatoolkit does not match cudaPackages.cudatoolkit" =
       MPISupport && cudaSupport && (mpi.cudatoolkit != cudaPackages.cudatoolkit);
@@ -296,7 +311,7 @@
 in
 buildPythonPackage rec {
   pname = "torch";
-  version = "2.8.0";
+  version = "2.9.0-rc2";
   pyproject = true;

   stdenv = effectiveStdenv;
@@ -315,14 +330,21 @@
     repo = "pytorch";
     tag = "v${version}";
     fetchSubmodules = true;
-    hash = "sha256-5JDYFoBe6bC9Dz143Bm/5OEOWsQxjctAR9fI4f6G2W8=";
+    hash = "sha256-s7m6DANNKanDxWLXI8sksNyLo1Y0xQgYy6/XJKKygaA=";
   };

   patches = [
     ./mkl-rpath.patch
+    # Include cstdint.h for uint8_t definition to fix gcc 14 compilation.
+    (fetchpatch {
+      name = "gloo-cstdint.diff";
+      url = "https://github.com/pytorch/gloo/commit/54cbae0d3a67fa890b4c3d9ee162b7860315e341.diff";
+      hash = "sha256-SsNN7wLhfpGgsdwZ+cS36tNLf8SKpMlJK6ya8y4AYnk=";
+      stripLen = 1;
+      extraPrefix = "third_party/gloo/";
+    })
   ]
   ++ lib.optionals cudaSupport [ ./fix-cmake-cuda-toolkit.patch ]
-  ++ lib.optionals rocmSupport [ ./cmake-load-hip-invalid-state.diff ]
   ++ lib.optionals (stdenv.hostPlatform.isDarwin && stdenv.hostPlatform.isx86_64) [
     # pthreadpool added support for Grand Central Dispatch in April
     # 2020. However, this relies on functionality (DISPATCH_APPLY_AUTO)
@@ -348,8 +370,8 @@
     in
     ''
       substituteInPlace pyproject.toml \
-        --replace-fail "setuptools>=62.3.0,<80.0" \
-                       "setuptools>=62.3.0"
+        --replace-fail "setuptools>=70.1.0,<80.0" \
+                       "setuptools>=70.1.0"

       substituteInPlace cmake/public/cuda.cmake \
         --replace-fail \
@@ -363,8 +385,8 @@

       # annotations (3.7), print_function (3.0), with_statement (2.6) are all supported
       sed -i -e "/from __future__ import/d" **.py
-      substituteInPlace third_party/NNPACK/CMakeLists.txt \
-        --replace-fail "PYTHONPATH=" 'PYTHONPATH=$ENV{PYTHONPATH}:'
+      #substituteInPlace third_party/NNPACK/CMakeLists.txt \
+      #  --replace-fail "PYTHONPATH=" 'PYTHONPATH=$ENV{PYTHONPATH}:'
       # flag from cmakeFlags doesn't work, not clear why
       # setting it at the top of NNPACK's own CMakeLists does
       sed -i '2s;^;set(PYTHON_SIX_SOURCE_DIR ${six.src})\n;' third_party/NNPACK/CMakeLists.txt
@@ -374,9 +396,9 @@
         --replace-fail 'addr2line_binary_ = "addr2line"' 'addr2line_binary_ = "${lib.getExe' binutils "addr2line"}"'

       # gen_pyi needs typing-extensions.
-      substituteInPlace torch/CMakeLists.txt \
-        --replace-fail "env PYTHONPATH=\"\''${TORCH_ROOT}\"" \
-                       "env PYTHONPATH=\"\''${TORCH_ROOT}:${pyiGenPath}\""
+      #substituteInPlace torch/CMakeLists.txt \
+      #  --replace-fail "env PYTHONPATH=\"\''${TORCH_ROOT}\"" \
+      #                 "env PYTHONPATH=\"\''${TORCH_ROOT}:${pyiGenPath}\""
     ''
     + lib.optionalString rocmSupport ''
       # https://github.com/facebookincubator/gloo/pull/297
@@ -386,11 +408,6 @@
       # Replace hard-coded rocm paths
       substituteInPlace caffe2/CMakeLists.txt \
         --replace-fail "/opt/rocm" "${rocmtoolkit_joined}"
-
-      # Strangely, this is never set in cmake
-      substituteInPlace cmake/public/LoadHIP.cmake \
-        --replace-fail "set(ROCM_PATH \$ENV{ROCM_PATH})" \
-          "set(ROCM_PATH \$ENV{ROCM_PATH})''\nset(ROCM_VERSION ${lib.concatStrings (lib.intersperse "0" (lib.splitVersion rocmPackages.clr.version))})"
     ''
     # Detection of NCCL version doesn't work particularly well when using the static binary.
     + lib.optionalString cudaSupport ''