py-onnxruntime: add v1.18.0 -> v1.18.3 and add ROCm support (#46448)

* add ROCm support for py-onnxruntime

* add new versions of py-onnxruntime

* add review changes

2 files changed, 115 insertions, 1 deletions

diff --git a/var/spack/repos/builtin/packages/py-onnxruntime/0001-Find-ROCm-Packages-Individually.patch b/var/spack/repos/builtin/packages/py-onnxruntime/0001-Find-ROCm-Packages-Individually.patch
new file mode 100644
index 0000000000..a48c192182
--- /dev/null
+++ b/var/spack/repos/builtin/packages/py-onnxruntime/0001-Find-ROCm-Packages-Individually.patch
@@ -0,0 +1,45 @@
+From c4930c939cc1c8b4c6122b1e9530942ecd517fb2 Mon Sep 17 00:00:00 2001
+From: Afzal Patel <Afzal.Patel@amd.com>
+Date: Tue, 17 Sep 2024 19:33:51 +0000
+Subject: [PATCH] Find individual ROCm dependencies
+
+---
+ cmake/onnxruntime_providers_rocm.cmake | 15 ++++++++++++++-
+ 1 file changed, 14 insertions(+), 1 deletion(-)
+
+diff --git a/cmake/onnxruntime_providers_rocm.cmake b/cmake/onnxruntime_providers_rocm.cmake
+index b662682915..2e9574c04d 100644
+--- a/cmake/onnxruntime_providers_rocm.cmake
++++ b/cmake/onnxruntime_providers_rocm.cmake
+@@ -11,6 +11,12 @@
+   find_package(rocblas REQUIRED)
+   find_package(MIOpen REQUIRED)
+   find_package(hipfft REQUIRED)
++  find_package(rocrand REQUIRED)
++  find_package(hipsparse REQUIRED)
++  find_package(hipcub REQUIRED)
++  find_package(rocprim REQUIRED)
++  find_package(rocthrust REQUIRED)
++  find_package(hipblas REQUIRED)
+
+   # MIOpen version
+   if(NOT DEFINED ENV{MIOPEN_PATH})
+@@ -147,7 +153,14 @@
+       ${eigen_INCLUDE_DIRS}
+     PUBLIC
+       ${onnxruntime_ROCM_HOME}/include
+-      ${onnxruntime_ROCM_HOME}/include/roctracer)
++      ${onnxruntime_ROCM_HOME}/include/roctracer
++      ${HIPRAND_INCLUDE_DIR}
++      ${ROCRAND_INCLUDE_DIR}
++      ${HIPSPARSE_INCLUDE_DIR}
++      ${HIPCUB_INCLUDE_DIR}
++      ${ROCPRIM_INCLUDE_DIR}
++      ${ROCTHRUST_INCLUDE_DIR}
++      ${HIPBLAS_INCLUDE_DIR})
+
+   set_target_properties(onnxruntime_providers_rocm PROPERTIES LINKER_LANGUAGE CXX)
+   set_target_properties(onnxruntime_providers_rocm PROPERTIES FOLDER "ONNXRuntime")
+--
+2.43.5
+
diff --git a/var/spack/repos/builtin/packages/py-onnxruntime/package.py b/var/spack/repos/builtin/packages/py-onnxruntime/package.py
index 14d1f54715..7ea80f69d6 100644
--- a/var/spack/repos/builtin/packages/py-onnxruntime/package.py
+++ b/var/spack/repos/builtin/packages/py-onnxruntime/package.py
@@ -6,7 +6,7 @@
 from spack.package import *
 
 
-class PyOnnxruntime(CMakePackage, PythonExtension):
+class PyOnnxruntime(CMakePackage, PythonExtension, ROCmPackage):
     """ONNX Runtime is a performance-focused complete scoring
     engine for Open Neural Network Exchange (ONNX) models, with
     an open extensible architecture to continually address the
@@ -22,6 +22,9 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
 
     license("MIT")
 
+    version("1.18.2", tag="v1.18.2", commit="9691af1a2a17b12af04652f4d8d2a18ce9507025")
+    version("1.18.1", tag="v1.18.1", commit="387127404e6c1d84b3468c387d864877ed1c67fe")
+    version("1.18.0", tag="v1.18.0", commit="45737400a2f3015c11f005ed7603611eaed306a6")
     version("1.17.3", tag="v1.17.3", commit="56b660f36940a919295e6f1e18ad3a9a93a10bf7")
     version("1.17.1", tag="v1.17.1", commit="8f5c79cb63f09ef1302e85081093a3fe4da1bc7d")
     version("1.10.0", tag="v1.10.0", commit="0d9030e79888d1d5828730b254fedc53c7b640c1")
@@ -50,6 +53,8 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
     depends_on("py-coloredlogs", when="@1.17:", type=("build", "run"))
     depends_on("py-flatbuffers", type=("build", "run"))
     depends_on("py-numpy@1.16.6:", type=("build", "run"))
+    depends_on("py-numpy@1.21.6:", when="@1.18:", type=("build", "run"))
+    depends_on("py-numpy@:1", when="@:1.18", type=("build", "run"))
     depends_on("py-packaging", type=("build", "run"))
     depends_on("py-protobuf", type=("build", "run"))
     depends_on("py-sympy@1.1:", type=("build", "run"))
@@ -60,6 +65,7 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
     depends_on("py-cerberus", type=("build", "run"))
     depends_on("py-onnx", type=("build", "run"))
     depends_on("py-onnx@:1.15.0", type=("build", "run"), when="@:1.17")
+    depends_on("py-onnx@:1.16", type=("build", "run"), when="@:1.18")
     depends_on("zlib-api")
     depends_on("libpng")
     depends_on("cuda", when="+cuda")
@@ -67,6 +73,35 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
     depends_on("iconv", type=("build", "link", "run"))
     depends_on("re2+shared")
 
+    rocm_dependencies = [
+        "hsa-rocr-dev",
+        "hip",
+        "hiprand",
+        "hipsparse",
+        "hipfft",
+        "hipcub",
+        "hipblas",
+        "llvm-amdgpu",
+        "miopen-hip",
+        "migraphx",
+        "rocblas",
+        "rccl",
+        "rocprim",
+        "rocminfo",
+        "rocm-core",
+        "rocm-cmake",
+        "roctracer-dev",
+        "rocthrust",
+        "rocrand",
+        "rocsparse",
+    ]
+
+    with when("+rocm"):
+        for pkg_dep in rocm_dependencies:
+            depends_on(f"{pkg_dep}@5.7:6.1", when="@1.17")
+            depends_on(f"{pkg_dep}@6.1:", when="@1.18:")
+            depends_on(pkg_dep)
+
     # Adopted from CMS experiment's fork of onnxruntime
     # https://github.com/cms-externals/onnxruntime/compare/5bc92df...d594f80
     patch("cms.patch", level=1, when="@1.7.2")
@@ -85,6 +120,10 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
         when="@1.10:1.15",
     )
 
+    # ORT is assuming all ROCm components are installed in a single path,
+    # this patch finds the packages individually
+    patch("0001-Find-ROCm-Packages-Individually.patch", when="@1.17: +rocm")
+
     dynamic_cpu_arch_values = ("NOAVX", "AVX", "AVX2", "AVX512")
 
     variant(
@@ -99,10 +138,28 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
     root_cmakelists_dir = "cmake"
     build_directory = "."
 
+    def patch(self):
+        if self.spec.satisfies("@1.17 +rocm"):
+            filter_file(
+                r"${onnxruntime_ROCM_HOME}/.info/version-dev",
+                "{0}/.info/version".format(self.spec["rocm-core"].prefix),
+                "cmake/CMakeLists.txt",
+                string=True,
+            )
+        if self.spec.satisfies("@1.18: +rocm"):
+            filter_file(
+                r"${onnxruntime_ROCM_HOME}/.info/version",
+                "{0}/.info/version".format(self.spec["rocm-core"].prefix),
+                "cmake/CMakeLists.txt",
+                string=True,
+            )
+
     def setup_build_environment(self, env):
         value = self.spec.variants["dynamic_cpu_arch"].value
         value = self.dynamic_cpu_arch_values.index(value)
         env.set("MLAS_DYNAMIC_CPU_ARCH", str(value))
+        if self.spec.satisfies("+rocm"):
+            env.set("MIOPEN_PATH", self.spec["miopen-hip"].prefix)
 
     def setup_run_environment(self, env):
         value = self.spec.variants["dynamic_cpu_arch"].value
@@ -137,6 +194,18 @@ class PyOnnxruntime(CMakePackage, PythonExtension):
                 )
             )
 
+        if self.spec.satisfies("+rocm"):
+            args.extend(
+                (
+                    define("CMAKE_HIP_COMPILER", f"{self.spec['llvm-amdgpu'].prefix}/bin/clang++"),
+                    define("onnxruntime_USE_MIGRAPHX", "ON"),
+                    define("onnxruntime_MIGRAPHX_HOME", self.spec["migraphx"].prefix),
+                    define("onnxruntime_USE_ROCM", "ON"),
+                    define("onnxruntime_ROCM_HOME", self.spec["hip"].prefix),
+                    define("onnxruntime_ROCM_VERSION", self.spec["hip"].version),
+                    define("onnxruntime_USE_COMPOSABLE_KERNEL", "OFF"),
+                )
+            )
         return args
 
     @run_after("install")