Py-torch 2.9, 2.11 and ci build on rocm

.ci/gitlab/configs/linux/ci.yaml

@@ -13,6 +13,8 @@ ci:
       - wrf
       build-job:
         tags: [ "spack", "huge" ]
+        # 12h for long ROCm/ML rebuilds (maximum allowed for AWS runners)
+        timeout: 720 minutes
         variables:
           CI_JOB_SIZE: huge
           SPACK_BUILD_JOBS: "12"

repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.5.patch

@@ -25,17 +25,18 @@ index 9be7f37..39d0f24 100644
  endif()
 
 diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake
-index 1c0d3a2..e0de4b1 100644
+index 1c0d3a2..83f9f9d 100644
 --- a/cmake/public/LoadHIP.cmake
 +++ b/cmake/public/LoadHIP.cmake
-@@ -167,6 +167,10 @@ if(HIP_FOUND)
+@@ -167,6 +167,11 @@ if(HIP_FOUND)
    find_package_and_print_version(hipsolver REQUIRED)
    find_package_and_print_version(hiprtc REQUIRED)
 
 +  list(APPEND ROCM_INCLUDE ${rocthrust_INCLUDE_DIR})
 +  list(APPEND ROCM_INCLUDE ${rocprim_INCLUDE_DIR})
 +  list(APPEND ROCM_INCLUDE ${hipcub_INCLUDE_DIR})
 +  list(APPEND ROCM_INCLUDE ${rocRAND_INCLUDE_DIR})
++  list(APPEND ROCM_INCLUDE $ENV{AOTRITON_INSTALLED_PREFIX}/include)
 
    find_library(PYTORCH_HIP_LIBRARIES amdhip64 HINTS ${ROCM_PATH}/lib)
    # TODO: miopen_LIBRARIES should return fullpath to the library file,

repos/spack_repo/builtin/packages/py_torch/PR152569-Update-spack-includes-2.7.patch

@@ -1,5 +1,5 @@
 diff --git a/caffe2/CMakeLists.txt b/caffe2/CMakeLists.txt
-index d2d23b7..620a89f 100644
+index d2d23b7ab65..620a89f65cb 100644
 --- a/caffe2/CMakeLists.txt
 +++ b/caffe2/CMakeLists.txt
 @@ -1379,13 +1379,6 @@ if(USE_ROCM)
@@ -26,7 +26,7 @@ index d2d23b7..620a89f 100644
  endif()
 
 diff --git a/cmake/public/LoadHIP.cmake b/cmake/public/LoadHIP.cmake
-index 58c74dd..d3e1ad4 100644
+index 58c74ddda35..54f96871372 100644
 --- a/cmake/public/LoadHIP.cmake
 +++ b/cmake/public/LoadHIP.cmake
 @@ -26,12 +26,6 @@ else()
@@ -78,7 +78,15 @@ index 58c74dd..d3e1ad4 100644
    find_package_and_print_version(amd_comgr REQUIRED)
    find_package_and_print_version(rocrand REQUIRED)
    find_package_and_print_version(hiprand REQUIRED)
-@@ -171,7 +168,11 @@ if(HIP_FOUND)
+@@ -157,6 +154,7 @@ if(HIP_FOUND)
+   find_package_and_print_version(hipcub REQUIRED)
+   find_package_and_print_version(rocthrust REQUIRED)
+   find_package_and_print_version(hipsolver REQUIRED)
++  list(APPEND ROCM_INCLUDE_DIRS $ENV{AOTRITON_INSTALLED_PREFIX}/include)
+   # workaround cmake 4 build issue
+   if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0")
+     message(WARNING "Work around hiprtc cmake failure for cmake >= 4")
+@@ -171,7 +169,11 @@ if(HIP_FOUND)
    if(UNIX)
      find_package_and_print_version(rccl)
      find_package_and_print_version(hsa-runtime64 REQUIRED)
@@ -89,4 +97,3 @@ index 58c74dd..d3e1ad4 100644
 +  if(UNIX)
      # roctx is part of roctracer
      find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib)
-

repos/spack_repo/builtin/packages/py_torch/package.py

@@ -97,7 +97,16 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     variant("rocm", default=False, description="Use ROCm")
     variant("cudnn", default=not is_darwin, description="Use cuDNN", when="+cuda")
     variant("fbgemm", default=True, description="Use FBGEMM (quantized 8-bit server operators)")
-    variant("kineto", default=True, description="Use Kineto profiling library", when="@1.8:")
+    variant("kineto", default=True, description="Use Kineto profiling library", when="@1.8:2.10")
+    variant(
+        "kineto", default=True, description="Use Kineto profiling library", when="@2.11: ~rocm"
+    )
+    variant(
+        "kineto",
+        default=False,
+        description="Disable Kineto from 2.11 on ROCm",
+        when="@2.11: +rocm",
+    )
     variant("magma", default=not is_darwin, description="Use MAGMA", when="+cuda")
     variant("metal", default=is_darwin, description="Use Metal for Caffe2 iOS build")
     variant(
@@ -107,7 +116,13 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
         when="@1.12: platform=darwin",
     )
     variant("nccl", default=True, description="Use NCCL", when="+cuda platform=linux")
-    variant("nccl", default=True, description="Use NCCL", when="+rocm platform=linux")
+    variant("nccl", default=True, description="Use NCCL", when="@:2.10 +rocm platform=linux")
+    variant(
+        "nccl",
+        default=False,
+        description="Disable NCCL from 2.11 on rocm",
+        when="@2.11: +rocm platform=linux",
+    )
     # Requires AVX2: https://discuss.pytorch.org/t/107518
     variant("nnpack", default=True, description="Use NNPACK", when="target=x86_64_v3:")
     variant("numa", default=True, description="Use NUMA", when="platform=linux")
@@ -144,6 +159,9 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     conflicts("+gloo+rocm")
     conflicts("+rocm", when="@2.3", msg="Rocm doesn't support py-torch 2.3 release")
     conflicts("+rocm", when="@2.4", msg="Rocm doesn't support py-torch 2.4 release")
+    conflicts("+rocm", when="@2.8", msg="Rocm doesn't support py-torch 2.8 release")
+    conflicts("+rocm", when="@2.12", msg="Rocm doesn't support py-torch 2.12 release")
+    conflicts("+rocm", when="@2.13", msg="Rocm doesn't support py-torch 2.13 release")
     conflicts("+tensorpipe", when="+rocm ^hip@:5.1", msg="TensorPipe not supported until ROCm 5.2")
     conflicts("+breakpad", when="target=ppc64:")
     conflicts("+breakpad", when="target=ppc64le:")
@@ -344,7 +362,8 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
     depends_on("valgrind", when="+valgrind")
     with when("+rocm"):
         depends_on("hsa-rocr-dev")
-        depends_on("hip")
+        depends_on("hip@7.0:", when="@2.9:")
+        depends_on("hip@:6.4", when="@:2.7")
         depends_on("rccl", when="+nccl")
         depends_on("rocprim")
         depends_on("hipcub")
@@ -356,14 +375,22 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
         depends_on("hiprand")
         depends_on("hipsolver")
         depends_on("rocm-core")
+        depends_on("amdsmi", when="@2.12")
         depends_on("rocfft")
         depends_on("rocblas")
         depends_on("miopen-hip")
+        depends_on("composable-kernel")
+        depends_on("hipblaslt")
+        depends_on("rocm-smi-lib")
+        # Ensure hipblaslt version for 2.9+
+        depends_on("hipblaslt@7.0:", when="@2.9:")
         depends_on("rocminfo")
-        depends_on("aotriton@0.8.1b", when="@2.5:2.6")
-        depends_on("aotriton@0.9.1b", when="@2.7:")
-        depends_on("composable-kernel@:6.3.2", when="@2.5")
-        depends_on("composable-kernel@6.3.2:", when="@2.6:")
+        depends_on("hipsparselt@7.0:", when="@2.9:")
+        depends_on("aotriton@0.8b", when="@2.5:2.6")
+        depends_on("aotriton@0.9.2b", when="@2.7")
+        depends_on("aotriton@0.10b", when="@2.8:2.10")
+        depends_on("aotriton@0.11b", when="@2.11:")
+
     depends_on("mpi", when="+mpi")
     depends_on("ucc", when="+ucc")
     depends_on("ucx", when="+ucc")
@@ -627,6 +654,17 @@ def patch(self):
             "torch_global_deps PROPERTIES LINKER_LANGUAGE CXX",
             "caffe2/CMakeLists.txt",
         )
+        if self.spec.satisfies("@2.5:+rocm"):
+            filter_file(
+                "find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib)",
+                "find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib)\n"
+                "if(DEFINED ENV{ROCTRACER_INCLUDE_DIR})\n"
+                "  set(ROCTRACER_INCLUDE_DIR $ENV{ROCTRACER_INCLUDE_DIR} CACHE PATH "
+                '"Roctracer include directory" FORCE)\n'
+                "endif()",
+                "cmake/public/LoadHIP.cmake",
+                string=True,
+            )
         if self.spec.satisfies("@2.1:2.7+rocm"):
             filter_file(
                 "${ROCM_INCLUDE_DIRS}/rocm-core/rocm_version.h",
@@ -724,6 +762,19 @@ def enable_or_disable(variant, keyword="USE", var=None):
 
         enable_or_disable("rocm")
         if "+rocm" in self.spec:
+            # So libtorch_hip.so and dependents find ROCm/runtime libs at runtime and
+            # during binary cache relocation (avoids "=> not found" for e.g.
+            # libamdhip64.so.6, libhsa-runtime64.so.1).
+            for lib_dir in [
+                self.spec["hip"].prefix.lib,
+                self.spec["hsa-rocr-dev"].prefix.lib,
+                self.spec["rocm-smi-lib"].prefix.lib,
+            ]:
+                env.append_flags("LDFLAGS", "-Wl,-rpath," + lib_dir)
+            # Link to rocm-smi-lib which provides rsmi_* symbols used by libtorch_hip.so
+            env.append_flags(
+                "LDFLAGS", "-L{} -lrocm_smi64".format(self.spec["rocm-smi-lib"].prefix.lib)
+            )
             env.set("PYTORCH_ROCM_ARCH", ";".join(self.spec.variants["amdgpu_target"].value))
             env.set("HSA_PATH", self.spec["hsa-rocr-dev"].prefix)
             env.set("ROCBLAS_PATH", self.spec["rocblas"].prefix)
@@ -744,6 +795,7 @@ def enable_or_disable(variant, keyword="USE", var=None):
             if self.spec.satisfies("@2.5:"):
                 env.set("TORCHINDUCTOR_CK_DIR", self.spec["composable-kernel"].prefix)
                 env.set("AOTRITON_INSTALLED_PREFIX", self.spec["aotriton"].prefix)
+                env.prepend_path("CPATH", self.spec["aotriton"].prefix.include)
             if self.spec.satisfies("^hip@5.2.0:"):
                 env.set("CMAKE_MODULE_PATH", self.spec["hip"].prefix.lib.cmake.hip)
 
@@ -867,6 +919,22 @@ def enable_or_disable(variant, keyword="USE", var=None):
 
     def setup_run_environment(self, env: EnvironmentModifications) -> None:
         self.torch_cuda_arch_list(env)
+        if "+rocm" in self.spec:
+            env.prepend_path("LD_LIBRARY_PATH", self.spec["hip"].prefix.lib)
+
+    def setup_dependent_build_environment(self, env, dependent_spec):
+        if "+rocm" in self.spec:
+            env.prepend_path("LD_LIBRARY_PATH", self.spec["hip"].prefix.lib)
+            # PyTorch headers (e.g. c10/util/complex.h) include <thrust/complex.h>;
+            # dependents need rocthrust include so HIP device builds can find it.
+            env.set("THRUST_PATH", self.spec["rocthrust"].prefix)
+            env.prepend_path("CPATH", self.spec["rocthrust"].prefix.include)
+
+    def setup_dependent_run_environment(self, env, dependent_spec):
+        """So dependents (e.g. py-torch-nvidia-apex, py-torchaudio) can find
+        libamdhip64.so when importing torch or running code that uses ROCm."""
+        if "+rocm" in self.spec:
+            env.prepend_path("LD_LIBRARY_PATH", self.spec["hip"].prefix.lib)
 
     @run_before("install")
     def build_amd(self):

repos/spack_repo/builtin/packages/py_torchvision/package.py

@@ -61,6 +61,9 @@ class PyTorchvision(PythonPackage):
     version("0.9.0", sha256="9351ed92aded632f8c7f59dfadac13c191a834babe682f5785ea47e6fcf6b472")
     version("0.8.2", sha256="9a866c3c8feb23b3221ce261e6153fc65a98ce9ceaa71ccad017016945c178bf")
 
+    # Fix duplicate symbol error when building with ROCm
+    patch("torchvision-0.26.0-rocm-vision-duplicate-symbol.patch", when="@0.26.0 ^py-torch+rocm")
+
     desc = "Enable support for native encoding/decoding of {} formats in torchvision.io"
     variant("png", default=True, description=desc.format("PNG"))
     variant("jpeg", default=True, description=desc.format("JPEG"))
@@ -207,6 +210,28 @@ def setup_build_environment(self, env: EnvironmentModifications) -> None:
             include.extend(query.headers.directories)
             library.extend(query.libs.directories)
 
+        # When building with ROCm, add all ROCm library include paths for HIP compilation
+        # PyTorch headers transitively include many ROCm headers that extensions need
+        if "^py-torch+rocm" in self.spec:
+            rocm_deps = [
+                "rocthrust",
+                "rocprim",
+                "hipsparse",
+                "hipblas",
+                "hipblas-common",
+                "hipblaslt",
+                "hipfft",
+                "hiprand",
+                "hipsolver",
+                "rocblas",
+                "rocsparse",
+                "rocsolver",
+                "rocfft",
+            ]
+            for dep in rocm_deps:
+                if dep in self.spec:
+                    include.append(self.spec[dep].prefix.include)
+
         # CONTRIBUTING.md says to use TORCHVISION_INCLUDE and TORCHVISION_LIBRARY, but
         # these do not work for older releases. Build uses a mix of Spack's compiler wrapper
         # and the actual compiler, so this is needed to get parts of the build working.
@@ -215,3 +240,10 @@ def setup_build_environment(self, env: EnvironmentModifications) -> None:
         env.set("TORCHVISION_LIBRARY", ":".join(library))
         env.set("CPATH", ":".join(include))
         env.set("LIBRARY_PATH", ":".join(library))
+
+        # For ROCm builds, also prepend ROCm includes to ensure hipcc can find them
+        if "^py-torch+rocm" in self.spec:
+            for dep in rocm_deps:
+                if dep in self.spec:
+                    env.prepend_path("CPATH", self.spec[dep].prefix.include)
+                    env.prepend_path("CPLUS_INCLUDE_PATH", self.spec[dep].prefix.include)

repos/spack_repo/builtin/packages/py_torchvision/torchvision-0.26.0-rocm-vision-duplicate-symbol.patch

@@ -0,0 +1,23 @@
+diff --git a/setup.py b/setup.py
+index 7bf83fc..a9c8e5d 100644
+--- a/setup.py
++++ b/setup.py
+@@ -153,6 +153,10 @@ def make_C_extension():
+         + list(CSRS_DIR.glob("ops/quantized/cpu/*.cpp"))
+     )
+     mps_sources = list(CSRS_DIR.glob("ops/mps/*.mm"))
++
++    # Exclude vision_hip.cpp from main sources when building with ROCm to avoid duplicate symbols
++    if IS_ROCM:
++        sources = [s for s in sources if "vision_hip.cpp" not in str(s)]
+
+     if IS_ROCM:
+         from torch.utils.hipify import hipify_python
+@@ -161,6 +165,7 @@ def make_C_extension():
+             project_directory=str(ROOT_DIR),
+             output_directory=str(ROOT_DIR),
+             includes="torchvision/csrc/ops/cuda/*",
++            ignores=("torchvision/csrc/vision.cpp",),
+             show_detailed=True,
+             is_pytorch_extension=True,
+         )

stacks/ml-linux-x86_64/spack.yaml

@@ -37,7 +37,6 @@ spack:
     - py-timm
     - py-torch
     - py-torch-geometric
-    - py-torch-nvidia-apex
     - py-torchaudio
     - py-torchgeo
     - py-torchmetrics
@@ -84,6 +83,7 @@ spack:
     - $scikit_learn_specs
     - $tensorboard_specs
     - $tensorflow_specs
+    - py-torch-nvidia-apex
     # xgboost requires older cuda, which doesn't support newer GCC
     # - $xgboost_specs
     override:
@@ -113,15 +113,11 @@ spack:
     - $jax_specs
     - matrix:
       - [$keras_specs]
-      exclude:
-      - py-keras backend=torch
-    # Does not yet support Spack-installed ROCm
-    # - $torch_specs
+    - $torch_specs
     - $scikit_learn_specs
     - $tensorboard_specs
     - $tensorflow_specs
-    # Does not yet support Spack-installed ROCm
-    # - $xgboost_specs
+    - $xgboost_specs
     override:
       packages:
         all:
@@ -139,9 +135,13 @@ spack:
             - target=x86_64_v3
             - ~cuda
             - +rocm
+            - ~distributed
             - amdgpu_target=gfx90a
-            - ~flash_attention
+            - +flash_attention
             - "%c,cxx=gcc"
+        hwloc:
+          require:
+            - ~rocm
 
   ci:
     pipeline-gen: