Enabling py-torch 2.9, 2.10 and ci build on rocm

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

@@ -13,6 +13,8 @@ ci:
       - wrf
       build-job:
         tags: [ "spack", "huge" ]
+        # 24h for long ROCm/ML rebuilds; GitLab project/runner max must allow this
+        timeout: 1440 minutes
         variables:
           CI_JOB_SIZE: huge
           SPACK_BUILD_JOBS: "12"

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

@@ -60,7 +60,7 @@ class Aotriton(CMakePackage):
     depends_on("pkgconfig", type="build")
 
     # build llvm version with mlir with the commit that matches inside the llvm-hash.txt
-    depends_on("aotriton-llvm@0.10", when="@0.10b")
+    depends_on("aotriton-llvm@0.10", when="@0.10b:")
     depends_on("aotriton-llvm@0.9", when="@0.9b")
     depends_on("aotriton-llvm@0.8", when="@0.8b")
 
@@ -134,4 +134,21 @@ def cmake_args(self):
         args = []
         args.append(self.define("AOTRITON_GPU_BUILD_TIMEOUT", 0))
         args.append(self.define("AOTRITON_NOIMAGE_MODE", "ON"))
+        # So libaotriton_v2.so and extensions find libamdhip64.so at runtime and
+        # during binary cache relocation (avoids "libamdhip64.so.6 => not found").
+        args.append(self.define("CMAKE_INSTALL_RPATH", self.spec["hip"].prefix.lib))
+        args.append(self.define("CMAKE_INSTALL_RPATH_USE_LINK_PATH", True))
+        # So libaotriton_v2.so and extensions find shared libs at runtime and
+        # during binary cache relocation (avoids "=> not found" for e.g.
+        # libamdhip64.so.6, libz.so.1, libhsa-runtime64.so.1, libc++abi.so.1,
+        # libunwind.so.1).
+        rpath_dirs = [
+            self.spec["hip"].prefix.lib,
+            self.spec["hsa-rocr-dev"].prefix.lib,
+            self.spec["zlib-api"].prefix.lib,
+            self.spec["aotriton-llvm"].prefix.lib,
+        ]
+        args.append(self.define("CMAKE_INSTALL_RPATH", rpath_dirs))
+        if self.spec.satisfies("@0.11b"):
+            args.append(self.define("AOTRITON_USE_TORCH", "OFF"))
         return args

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

@@ -134,7 +134,9 @@ def cmake_args(self):
             args.append(self.define("INSTANCES_ONLY", "ON"))
         if self.run_tests:
             args.append(self.define("BUILD_TESTING", "ON"))
-        elif self.spec.satisfies("@:6.1"):
+        else:
+            args.append(self.define("BUILD_TESTING", "OFF"))
+        if self.spec.satisfies("@:6.1"):
             args.append(self.define("INSTANCES_ONLY", "ON"))
         if self.spec.satisfies("@:5.7"):
             args.append(self.define("CMAKE_CXX_FLAGS", "-O3"))

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

@@ -70,7 +70,7 @@ def url_for_version(self, version):
 
     depends_on("cmake@3.25.2:", type="build", when="@6.2.0:")
     depends_on("python@3.7:")
-    depends_on("python@3.8:3.13.2", when="@6.4:")
+    depends_on("python@3.8:3.13", when="@6.4:")
 
     for ver in [
         "6.0.0",

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

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
-import itertools
 import re
 
 from spack_repo.builtin.build_systems.cmake import CMakePackage
@@ -119,12 +118,12 @@ def url_for_version(self, version):
         depends_on(f"rocm-cmake@{ver}:", type="build", when=f"@{ver}")
         depends_on(f"hip@{ver}", when=f"@{ver}")
         depends_on(f"rocm-clang-ocl@{ver}", when=f"@{ver}")
-        depends_on(f"rocblas@{ver}", when=f"@{ver}")
-        for tgt in itertools.chain(["auto"], amdgpu_targets):
+        for tgt in ROCmPackage.amdgpu_targets:
             depends_on(
                 f"composable-kernel@{ver} amdgpu_target={tgt}",
                 when=f"@{ver} +ck amdgpu_target={tgt}",
             )
+            depends_on(f"rocblas@{ver} amdgpu_target={tgt}", when=f"@{ver} amdgpu_target={tgt}")
 
     for ver in ["6.0.0", "6.0.2", "6.1.0", "6.1.1", "6.1.2"]:
         depends_on(f"roctracer-dev@{ver}", when=f"@{ver}")

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

@@ -61,6 +61,10 @@ class PyLlvmlite(PythonPackage):
     depends_on("llvm@14", when="@0.41:0.43")
     depends_on("llvm@11:14", when="@0.40")
     depends_on("llvm@11", when="@0.37:0.39")
+    # When the spec has both llvm@20 and hwloc+rocm (e.g. llvm build-depends on hwloc),
+    # hwloc with older rocm versions conflicts with llvm@20 while linking.
+    # Hence adding hwloc~rocm dependency for llvmlite@0.45:.
+    depends_on("hwloc ~rocm", when="@0.45:")
     for t in [
         "arm:",
         "ppc:",

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)

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

@@ -141,6 +141,7 @@ 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("+tensorpipe", when="+rocm ^hip@:5.1", msg="TensorPipe not supported until ROCm 5.2")
     conflicts("+breakpad", when="target=ppc64:")
     conflicts("+breakpad", when="target=ppc64le:")
@@ -328,7 +329,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")
@@ -343,11 +345,17 @@ class PyTorch(PythonPackage, CudaPackage, ROCmPackage):
         depends_on("rocfft")
         depends_on("rocblas")
         depends_on("miopen-hip")
+        depends_on("composable-kernel")
+        depends_on("hipblaslt")
+        # 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")
@@ -591,6 +599,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",
@@ -688,6 +707,11 @@ 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]:
+                env.append_flags("LDFLAGS", "-Wl,-rpath," + lib_dir)
             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)
@@ -708,6 +732,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)
 
@@ -824,6 +849,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/rocm_validation_suite/package.py

@@ -76,7 +76,7 @@ class RocmValidationSuite(CMakePackage):
     depends_on("googletest")
     depends_on("doxygen", type="build")
     depends_on("libdrm", when="@6.4:")
-    depends_on("pciutils+shared", when="@6.4:")
+    depends_on("pciutils+shared", type="build", when="@6.4:")
 
     def setup_build_environment(self, env: EnvironmentModifications) -> None:
         spec = self.spec

stacks/ml-linux-x86_64/spack.yaml

@@ -28,21 +28,21 @@ spack:
     - py-keras backend=jax
     - py-keras backend=torch
   - torch_specs:
-    - py-botorch
-    - py-gpytorch
-    - py-kornia
-    - py-lightning
-    - py-pytorch-lightning
-    - py-segmentation-models-pytorch
-    - py-timm
+    #- py-botorch
+    #- py-gpytorch
+    #- py-kornia
+    #- py-lightning
+    #- py-pytorch-lightning
+    #- py-segmentation-models-pytorch
+    #- py-timm
     - py-torch
-    - py-torch-geometric
-    - py-torch-nvidia-apex
-    - py-torchaudio
-    - py-torchgeo
-    - py-torchmetrics
-    - py-torchvision
-    - py-vector-quantize-pytorch
+    #- py-torch-geometric
+    #- py-torch-nvidia-apex
+    #- py-torchaudio
+    #- py-torchgeo
+    #- py-torchmetrics
+    #- py-torchvision
+    #- py-vector-quantize-pytorch
   - scikit_learn_specs:
     - py-scikit-learn
   - tensorboard_specs:
@@ -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,6 +135,7 @@ spack:
             - target=x86_64_v3
             - ~cuda
             - +rocm
+            - ~distributed
             - amdgpu_target=gfx90a
             - ~flash_attention
             - "%c,cxx=gcc"