Update PyTorch ecosystem (#42394)

5 files changed, 39 insertions, 15 deletions

diff --git a/var/spack/repos/builtin/packages/pthreadpool/package.py b/var/spack/repos/builtin/packages/pthreadpool/package.py
index 7f95850d18..64f190e4d0 100644
--- a/var/spack/repos/builtin/packages/pthreadpool/package.py
+++ b/var/spack/repos/builtin/packages/pthreadpool/package.py
@@ -15,7 +15,8 @@ class Pthreadpool(CMakePackage):
     license("BSD-2-Clause")
 
     version("master", branch="master")
-    version("2021-04-13", commit="a134dd5d4cee80cce15db81a72e7f929d71dd413")  # py-torch@1.9
+    version("2023-08-29", commit="4fe0e1e183925bf8cfa6aae24237e724a96479b8")  # py-torch@2.2
+    version("2021-04-13", commit="a134dd5d4cee80cce15db81a72e7f929d71dd413")  # py-torch@1.9:2.1
     version("2020-10-05", commit="fa75e65a58a5c70c09c30d17a1fe1c1dff1093ae")  # py-torch@1.8
     version("2020-06-15", commit="029c88620802e1361ccf41d1970bd5b07fd6b7bb")  # py-torch@1.6:1.7
     version("2019-10-29", commit="d465747660ecf9ebbaddf8c3db37e4a13d0c9103")  # py-torch@1.5
@@ -35,8 +36,8 @@ class Pthreadpool(CMakePackage):
     )
     resource(
         name="googletest",
-        url="https://github.com/google/googletest/archive/release-1.10.0.zip",
-        sha256="94c634d499558a76fa649edb13721dce6e98fb1e7018dfaeba3cd7a083945e91",
+        url="https://github.com/google/googletest/archive/release-1.12.0.zip",
+        sha256="ce7366fe57eb49928311189cb0e40e0a8bf3d3682fca89af30d884c25e983786",
         destination="deps",
         placement="googletest",
     )
diff --git a/var/spack/repos/builtin/packages/py-torch/package.py b/var/spack/repos/builtin/packages/py-torch/package.py
index 6126b33111..e9040963cf 100644
--- a/var/spack/repos/builtin/packages/py-torch/package.py
+++ b/var/spack/repos/builtin/packages/py-torch/package.py
@@ -23,9 +23,10 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     # core libraries to ensure that the package was successfully installed.
     import_modules = ["torch", "torch.autograd", "torch.nn", "torch.utils"]
 
-    license("Intel")
+    license("BSD-3-Clause")
 
     version("main", branch="main")
+    version("2.2.0", tag="v2.2.0", commit="8ac9b20d4b090c213799e81acf48a55ea8d437d6")
     version("2.1.2", tag="v2.1.2", commit="a8e7c98cb95ff97bb30a728c6b2a1ce6bff946eb")
     version("2.1.1", tag="v2.1.1", commit="4c55dc50355d5e923642c59ad2a23d6ad54711e7")
     version("2.1.0", tag="v2.1.0", commit="7bcf7da3a268b435777fe87c7794c382f444e86d")
@@ -129,9 +130,9 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     )
 
     # Required dependencies
-    # See python_min_version in setup.py
-    # Upper bounds come from wheel availability on PyPI
-    depends_on("python@3.8:3.11", when="@2:", type=("build", "link", "run"))
+    # Based on PyPI wheel availability
+    depends_on("python@3.8:3.12", when="@2.2:", type=("build", "link", "run"))
+    depends_on("python@3.8:3.11", when="@2.0:2.1", type=("build", "link", "run"))
     depends_on("python@:3.10", when="@1.11:1", type=("build", "link", "run"))
     depends_on("python@:3.9", when="@1.7.1:1.10", type=("build", "link", "run"))
     depends_on("python@:3.8", when="@1.4:1.7.0", type=("build", "link", "run"))
@@ -155,6 +156,7 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
 
     # setup.py
     depends_on("py-filelock", when="@2:", type=("build", "run"))
+    depends_on("py-typing-extensions@4.8:", when="@2.2:", type=("build", "run"))
     depends_on("py-typing-extensions@3.6.2.1:", when="@1.7:", type=("build", "run"))
     depends_on("py-sympy", when="@2:", type=("build", "run"))
     depends_on("py-networkx", when="@2:", type=("build", "run"))
@@ -190,7 +192,8 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     # depends_on("sleef@3.5.1_2020-12-22", when="@1.8:")
     # depends_on("sleef@3.4.0_2019-07-30", when="@1.6:1.7")
     depends_on("fp16@2020-05-14", when="@1.6:")
-    depends_on("pthreadpool@2021-04-13", when="@1.9:")
+    depends_on("pthreadpool@2023-08-29", when="@2.2:")
+    depends_on("pthreadpool@2021-04-13", when="@1.9:2.1")
     depends_on("pthreadpool@2020-10-05", when="@1.8")
     depends_on("pthreadpool@2020-06-15", when="@1.6:1.7")
     depends_on("psimd@2020-05-17", when="@1.6:")
@@ -198,13 +201,15 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     depends_on("benchmark", when="@1.6:+test")
 
     # Optional dependencies
-    # https://github.com/pytorch/pytorch#prerequisites
+    # cmake/public/cuda.cmake
     depends_on("cuda@11:", when="@2:+cuda", type=("build", "link", "run"))
     depends_on("cuda@10.2:", when="@1.11:1+cuda", type=("build", "link", "run"))
     # https://discuss.pytorch.org/t/compiling-1-10-1-from-source-with-gcc-11-and-cuda-11-5/140971
     depends_on("cuda@10.2:11.4", when="@1.10+cuda", type=("build", "link", "run"))
     depends_on("cuda@9.2:11.4", when="@1.6:1.9+cuda", type=("build", "link", "run"))
     depends_on("cuda@9:11.4", when="@:1.5+cuda", type=("build", "link", "run"))
+    # https://github.com/pytorch/pytorch#prerequisites
+    depends_on("cudnn@8.5:", when="@2.3:+cudnn")
     depends_on("cudnn@7:", when="@1.6:+cudnn")
     depends_on("cudnn@7", when="@:1.5+cudnn")
     depends_on("magma+cuda", when="+magma+cuda")
@@ -244,7 +249,8 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     depends_on("gloo@2020-03-17", when="@1.6+gloo")
     depends_on("gloo+cuda", when="@1.6:+gloo+cuda")
     # https://github.com/pytorch/pytorch/issues/60331
-    # depends_on("onnx@1.14.1", when="@2.1:+onnx_ml")
+    # depends_on("onnx@1.15.0", when="@2.2:+onnx_ml")
+    # depends_on("onnx@1.14.1", when="@2.1+onnx_ml")
     # depends_on("onnx@1.13.1", when="@2.0+onnx_ml")
     # depends_on("onnx@1.12.0", when="@1.13:1+onnx_ml")
     # depends_on("onnx@1.11.0", when="@1.12+onnx_ml")
@@ -259,6 +265,8 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     depends_on("py-six", type="test")
     depends_on("py-psutil", type="test")
 
+    conflicts("%gcc@:9.3", when="@2.2:", msg="C++17 support required")
+
     # https://github.com/pytorch/pytorch/issues/90448
     patch(
         "https://github.com/pytorch/pytorch/pull/97270.patch?full_index=1",
@@ -426,6 +434,13 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
         when="@:2.1",
     )
 
+    patch(
+        "https://github.com/pytorch/FBGEMM/commit/da01a59556fec9776733bf20aea8fe8fb29cdd3d.patch?full_index=1",
+        sha256="97d8bd43f8cd8bb203dab3480d609c08499224acaca9915f2bdeb23c62350fb1",
+        when="@2.0.1 +fbgemm",
+        working_dir="third_party/fbgemm",
+    )
+
     @when("@1.5.0:")
     def patch(self):
         # https://github.com/pytorch/pytorch/issues/52208
diff --git a/var/spack/repos/builtin/packages/py-torchaudio/package.py b/var/spack/repos/builtin/packages/py-torchaudio/package.py
index 9d0d39798b..4fd49f05e6 100644
--- a/var/spack/repos/builtin/packages/py-torchaudio/package.py
+++ b/var/spack/repos/builtin/packages/py-torchaudio/package.py
@@ -17,6 +17,7 @@ class PyTorchaudio(PythonPackage):
     license("BSD-2-Clause")
 
     version("main", branch="main")
+    version("2.2.0", tag="v2.2.0", commit="08901ade5d17d3e3cf6fc039cbd601cbd2853686")
     version("2.1.2", tag="v2.1.2", commit="c4c1957d24b423200fd83591d46066135979a5a8")
     version("2.1.1", tag="v2.1.1", commit="db624844f5c95bb7618fe5a5f532bf9b68efeb45")
     version("2.1.0", tag="v2.1.0", commit="6ea1133706801ec6e81bb29142da2e21a8583a0a")
@@ -42,8 +43,9 @@ class PyTorchaudio(PythonPackage):
     version("0.5.0", tag="v0.5.0", commit="09494ea545738538f9db2dceeffe10d421060ee5")
     version("0.4.0", tag="v0.4.0", commit="8afed303af3de41f3586007079c0534543c8f663")
 
-    # https://pytorch.org/audio/main/installation.html#dependencies
-    depends_on("python@3.8:3.11", when="@2:", type=("build", "link", "run"))
+    # Based on PyPI wheel availability
+    depends_on("python@3.8:3.12", when="@2.2:", type=("build", "link", "run"))
+    depends_on("python@3.8:3.11", when="@2.0:2.1", type=("build", "link", "run"))
     depends_on("python@:3.10", when="@0.12:0", type=("build", "link", "run"))
     depends_on("python@:3.9", when="@0.7.2:0.11", type=("build", "link", "run"))
     depends_on("python@:3.8", when="@:0.7.0", type=("build", "link", "run"))
@@ -60,6 +62,7 @@ class PyTorchaudio(PythonPackage):
     depends_on("sox")
 
     depends_on("py-torch@main", when="@main", type=("build", "link", "run"))
+    depends_on("py-torch@2.2.0", when="@2.2.0", type=("build", "link", "run"))
     depends_on("py-torch@2.1.2", when="@2.1.2", type=("build", "link", "run"))
     depends_on("py-torch@2.1.1", when="@2.1.1", type=("build", "link", "run"))
     depends_on("py-torch@2.1.0", when="@2.1.0", type=("build", "link", "run"))
diff --git a/var/spack/repos/builtin/packages/py-torchtext/package.py b/var/spack/repos/builtin/packages/py-torchtext/package.py
index 4758d1a2f8..a04fef47a3 100644
--- a/var/spack/repos/builtin/packages/py-torchtext/package.py
+++ b/var/spack/repos/builtin/packages/py-torchtext/package.py
@@ -19,6 +19,7 @@ class PyTorchtext(PythonPackage):
     license("BSD-3-Clause")
 
     version("main", branch="main")
+    version("0.17.0", tag="v0.17.0", commit="400da5c61bab4abaaeaeca91744ca031ad9b2edf")
     version("0.16.2", tag="v0.16.2", commit="299b90e908c1b492139a4cf9da3912660e79a06b")
     version("0.16.1", tag="v0.16.1", commit="66671007c84e07386da3c04e5ca403b8a417c8e5")
     version("0.16.0", tag="v0.16.0", commit="4e255c95c76b1ccde4f6650391c0bc30650d6dbe")
@@ -38,7 +39,7 @@ class PyTorchtext(PythonPackage):
     version("0.6.0", tag="0.6.0", commit="3a54c7f52584f201c17ca7489b52b812152612dc")
     version("0.5.0", tag="0.5.0", commit="0169cde2f1d446ae886ef0be07e9a673585ed256")
 
-    # https://github.com/pytorch/text#installation
+    # Based on PyPI wheel availability
     depends_on("python@3.8:3.11", when="@0.15:", type=("build", "link", "run"))
     depends_on("python@:3.10", when="@0.13:0.14", type=("build", "link", "run"))
     depends_on("python@:3.9", when="@0.8.1:0.12", type=("build", "link", "run"))
@@ -62,6 +63,7 @@ class PyTorchtext(PythonPackage):
 
     # https://github.com/pytorch/text#installation
     depends_on("py-torch@main", when="@main", type=("build", "link", "run"))
+    depends_on("py-torch@2.2.0", when="@0.17.0", type=("build", "link", "run"))
     depends_on("py-torch@2.1.2", when="@0.16.2", type=("build", "link", "run"))
     depends_on("py-torch@2.1.1", when="@0.16.1", type=("build", "link", "run"))
     depends_on("py-torch@2.1.0", when="@0.16.0", type=("build", "link", "run"))
diff --git a/var/spack/repos/builtin/packages/py-torchvision/package.py b/var/spack/repos/builtin/packages/py-torchvision/package.py
index dfb063cd35..ad97b8a64d 100644
--- a/var/spack/repos/builtin/packages/py-torchvision/package.py
+++ b/var/spack/repos/builtin/packages/py-torchvision/package.py
@@ -19,6 +19,7 @@ class PyTorchvision(PythonPackage):
     license("BSD-3-Clause")
 
     version("main", branch="main")
+    version("0.17.0", sha256="55e395d5c7d9bf7658c82ac633cac2224aa168e1bfe8bb5b2b2a296c792a3500")
     version("0.16.2", sha256="8c1f2951e98d8ada6e5a468f179af4be9f56d2ebc3ab057af873da61669806d7")
     version("0.16.1", sha256="d31fe52e4540750c8d372b0f38f1bfa81d8261193f2c2c06577332831d203c50")
     version("0.16.0", sha256="79b30b082237e3ead21e74587cedf4a4d832f977cf7dfeccfb65f67988b12ceb")
@@ -53,14 +54,16 @@ class PyTorchvision(PythonPackage):
     variant("ffmpeg", default=False, description=desc.format("FFMPEG"))
     variant("video_codec", default=False, description=desc.format("video_codec"))
 
-    # https://github.com/pytorch/vision#installation
-    depends_on("python@3.8:3.11", when="@0.15:", type=("build", "link", "run"))
+    # Based on PyPI wheel availability
+    depends_on("python@3.8:3.12", when="@0.17:", type=("build", "link", "run"))
+    depends_on("python@3.8:3.11", when="@0.15:0.16", type=("build", "link", "run"))
     depends_on("python@:3.10", when="@0.12:0.14", type=("build", "link", "run"))
     depends_on("python@:3.9", when="@0.8.2:0.11", type=("build", "link", "run"))
     depends_on("python@:3.8", when="@0.5:0.8.1", type=("build", "link", "run"))
 
     # https://github.com/pytorch/vision#installation
     depends_on("py-torch@main", when="@main", type=("build", "link", "run"))
+    depends_on("py-torch@2.2.0", when="@0.17.0", type=("build", "link", "run"))
     depends_on("py-torch@2.1.2", when="@0.16.2", type=("build", "link", "run"))
     depends_on("py-torch@2.1.1", when="@0.16.1", type=("build", "link", "run"))
     depends_on("py-torch@2.1.0", when="@0.16.0", type=("build", "link", "run"))