legion: add HIP support (#31080)

1 files changed, 33 insertions, 5 deletions

diff --git a/var/spack/repos/builtin/packages/legion/package.py b/var/spack/repos/builtin/packages/legion/package.py
index 087f19005b..35196106f8 100644
--- a/var/spack/repos/builtin/packages/legion/package.py
+++ b/var/spack/repos/builtin/packages/legion/package.py
@@ -8,7 +8,7 @@ import os
 from spack.package import *
 
 
-class Legion(CMakePackage):
+class Legion(CMakePackage, ROCmPackage):
     """Legion is a data-centric parallel programming system for writing
        portable high performance programs targeted at distributed heterogeneous
        architectures. Legion presents abstractions which allow programmers to
@@ -50,13 +50,29 @@ class Legion(CMakePackage):
     # TODO: we could use a map here to clean up and use naming vs. numbers.
     cuda_arch_list = ('60', '70', '75', '80')
     for nvarch in cuda_arch_list:
-        depends_on('kokkos@3.3.01+cuda+cuda_lambda+wrapper cuda_arch={0}'.format(nvarch),
+        depends_on('kokkos@3.3.01:+cuda+cuda_lambda+wrapper cuda_arch={0}'.format(nvarch),
                    when='%gcc+kokkos+cuda cuda_arch={0}'.format(nvarch))
-        depends_on("kokkos@3.3.01+cuda+cuda_lambda~wrapper cuda_arch={0}".format(nvarch),
+        depends_on("kokkos@3.3.01:+cuda+cuda_lambda~wrapper cuda_arch={0}".format(nvarch),
                    when="%clang+kokkos+cuda cuda_arch={0}".format(nvarch))
 
-    depends_on('kokkos@3.3.01~cuda', when='+kokkos~cuda')
-    depends_on("kokkos@3.3.01~cuda+openmp", when='+kokkos+openmp')
+    depends_on('kokkos@3.3.01:~cuda', when='+kokkos~cuda')
+    depends_on("kokkos@3.3.01:~cuda+openmp", when='+kokkos+openmp')
+
+    # HIP specific
+    variant('hip_hijack', default=False,
+            description="Hijack application calls into the HIP runtime",
+            when='+rocm')
+    variant('hip_target', default='ROCM',
+            values=('ROCM', 'CUDA'),
+            description="API used by HIP",
+            multi=False,
+            when='+rocm')
+
+    for arch in ROCmPackage.amdgpu_targets:
+        depends_on('kokkos@3.3.01:+rocm amdgpu_target={0}'.format(arch),
+                   when='+rocm amdgpu_target={0}'.format(arch))
+
+    depends_on('kokkos@3.3.01:+rocm', when='+kokkos+rocm')
 
     depends_on('python@3', when='+python')
     depends_on('papi', when='+papi')
@@ -205,6 +221,11 @@ class Legion(CMakePackage):
     variant('max_fields', values=int, default=512,
             description="Maximum number of fields allowed in a logical region.")
 
+    def setup_build_environment(self, build_env):
+        spec = self.spec
+        if '+rocm' in spec:
+            build_env.set("HIP_PATH", spec['hip'].prefix)
+
     def cmake_args(self):
         spec = self.spec
         cmake_cxx_flags = []
@@ -273,6 +294,13 @@ class Legion(CMakePackage):
             if '+cuda_unsupported_compiler' in spec:
                 options.append('-DCUDA_NVCC_FLAGS:STRING=--allow-unsupported-compiler')
 
+        if '+rocm' in spec:
+            options.append('-DLegion_USE_HIP=ON')
+            options.append('-DLegion_GPU_REDUCTIONS=ON')
+            options.append(from_variant('Legion_HIP_TARGET', 'hip_target'))
+            options.append(from_variant('Legion_HIP_ARCH', 'amdgpu_target'))
+            options.append(from_variant('Legion_HIJACK_HIP', 'hip_hijack'))
+
         if '+fortran' in spec:
             # default is off.
             options.append('-DLegion_USE_Fortran=ON')