Enabling ci build for py-torch 2.9, 2.10 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/hwloc/package.py

@@ -135,7 +135,8 @@ class Hwloc(AutotoolsPackage, CudaPackage, ROCmPackage):
     depends_on("mpi", when="+netloc")
 
     with when("+rocm"):
-        depends_on("rocm-smi-lib")
+        depends_on("rocm-smi-lib@:6.4", when="@:2.11.1")
+        depends_on("rocm-smi-lib@7.0:", when="@2.12.2:")
         depends_on("rocm-opencl", when="+opencl")
         # Avoid a circular dependency since the openmp
         # variant of llvm-amdgpu depends on hwloc.

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

@@ -123,6 +123,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:")
@@ -305,7 +306,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")
@@ -320,11 +322,15 @@ 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:")
     depends_on("mpi", when="+mpi")
     depends_on("ucc", when="+ucc")
     depends_on("ucx", when="+ucc")
@@ -568,6 +574,14 @@ 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"
+                "set(ROCTRACER_INCLUDE_DIR $ENV{ROCTRACER_INCLUDE_DIR})",
+                "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",

stacks/ml-linux-x86_64-rocm/spack.yaml

@@ -43,7 +43,7 @@ spack:
     # Keras
     - py-keras backend=tensorflow
     # - py-keras backend=jax
-    # - py-keras backend=torch
+    - py-keras backend=torch
 
     # PyTorch
     # Does not yet support Spack-installed ROCm
@@ -54,7 +54,7 @@ spack:
     # - py-pytorch-lightning
     # - py-segmentation-models-pytorch
     # - py-timm
-    # - py-torch
+    - py-torch
     # - py-torch-geometric
     # - py-torch-nvidia-apex
     # - py-torchaudio