Enable rocm support and gtest variant for ucx (#36693)

1 files changed, 35 insertions, 1 deletions

diff --git a/var/spack/repos/builtin/packages/ucx/package.py b/var/spack/repos/builtin/packages/ucx/package.py
index 2e9d6e906a..8686324e9e 100644
--- a/var/spack/repos/builtin/packages/ucx/package.py
+++ b/var/spack/repos/builtin/packages/ucx/package.py
@@ -115,6 +115,7 @@ class Ucx(AutotoolsPackage, CudaPackage):
     variant("ud", default=False, description="Compile with IB Unreliable Datagram support")
     variant("verbs", default=False, description="Build OpenFabrics support")
     variant("xpmem", default=False, description="Enable XPMEM support")
+    variant("gtest", default=False, description="Build and install Googletest")
 
     depends_on("binutils+ld", when="%aocc", type="build")
     depends_on("binutils", when="+backtrace_detail")
@@ -129,6 +130,7 @@ class Ucx(AutotoolsPackage, CudaPackage):
     depends_on("rdma-core", when="+rdmacm")
     depends_on("rdma-core", when="+verbs")
     depends_on("xpmem", when="+xpmem")
+    depends_on("hip", when="+rocm")
 
     conflicts("+gdrcopy", when="~cuda", msg="gdrcopy currently requires cuda support")
     conflicts("+rocm", when="+gdrcopy", msg="gdrcopy > 2.0 does not support rocm")
@@ -138,6 +140,19 @@ class Ucx(AutotoolsPackage, CudaPackage):
     # See https://github.com/openucx/ucx/pull/8629, wrong int type
     patch("commit-2523555.patch", when="@1.13.1")
 
+    def patch(self):
+        if self.spec.satisfies("+rocm"):
+            filter_file("$$with_rocm", "${with_rocm[@]}", "configure", string=True)
+            filter_file(
+                "-I$with_rocm/include/hip -I$with_rocm/include",
+                "$ROCM_CPPFLAGS",
+                "configure",
+                string=True,
+            )
+            filter_file(
+                "-L$with_rocm/hip/lib -L$with_rocm/lib", "$ROCM_LDFLAGS", "configure", string=True
+            )
+
     @when("@1.9-dev")
     def autoreconf(self, spec, prefix):
         Executable("./autogen.sh")()
@@ -155,10 +170,10 @@ class Ucx(AutotoolsPackage, CudaPackage):
         args += self.with_or_without("openmp")
         args += self.enable_or_disable("optimizations")
         args += self.enable_or_disable("params-check", variant="parameter_checking")
+        args += self.enable_or_disable("gtest")
         args += self.with_or_without("pic")
 
         args += self.with_or_without("cuda", activation_value="prefix")
-        args += self.with_or_without("rocm")  # todo, prefix, avoid /opt/rocm guess.
 
         args += self.with_or_without("cm")
         args += self.enable_or_disable("cma")
@@ -219,9 +234,28 @@ class Ucx(AutotoolsPackage, CudaPackage):
         if "%aocc" in spec:
             args.append("LDFLAGS=-fuse-ld=bfd")
 
+        if "+rocm" in spec:
+            rocm_flags = " ".join(
+                [
+                    "-I" + self.spec["hip"].prefix.include,
+                    "-I" + self.spec["hip"].prefix.include.hip,
+                    "-I" + self.spec["hsa-rocr-dev"].prefix.include.hsa,
+                    "-L" + self.spec["hip"].prefix.lib,
+                    "-L" + self.spec["hsa-rocr-dev"].prefix.lib,
+                ]
+            )
+            args.append("--with-rocm=" + rocm_flags)
+        else:
+            args.append("--without-rocm")
+
         return args
 
     @run_after("install")
     def drop_examples(self):
         if self.spec.satisfies("~examples"):
             shutil.rmtree(join_path(self.spec.prefix, "share", "ucx", "examples"))
+
+    @run_after("install")
+    def install_gtest(self):
+        if self.spec.satisfies("+gtest"):
+            install_tree("test", self.spec.prefix.test)