kripke: add cuda and rocm support (#34257)

1 files changed, 60 insertions, 12 deletions

diff --git a/var/spack/repos/builtin/packages/kripke/package.py b/var/spack/repos/builtin/packages/kripke/package.py
index 0fbdb7f59c..a549fa40a1 100644
--- a/var/spack/repos/builtin/packages/kripke/package.py
+++ b/var/spack/repos/builtin/packages/kripke/package.py
@@ -6,7 +6,7 @@
 from spack.package import *
 
 
-class Kripke(CMakePackage):
+class Kripke(CMakePackage, CudaPackage, ROCmPackage):
     """Kripke is a simple, scalable, 3D Sn deterministic particle
     transport proxy/mini app.
     """
@@ -15,32 +15,80 @@ class Kripke(CMakePackage):
     git = "https://github.com/LLNL/Kripke.git"
 
     tags = ["proxy-app"]
-    version("1.2.4", submodules=True, tag="v1.2.4")
+
+    maintainers = ["vsrana01"]
+
+    version("develop", branch="develop", submodules=False)
+    version("1.2.4", submodules=False, tag="v1.2.4")
     version("1.2.3", submodules=True, tag="v1.2.3")
     version("1.2.2", submodules=True, tag="v1.2.2-CORAL2")
     version("1.2.1", submodules=True, tag="v1.2.1-CORAL2")
     version("1.2.0", submodules=True, tag="v1.2.0-CORAL2")
 
     variant("mpi", default=True, description="Build with MPI.")
-    variant("openmp", default=True, description="Build with OpenMP enabled.")
+    variant("openmp", default=False, description="Build with OpenMP enabled.")
     variant("caliper", default=False, description="Build with Caliper support enabled.")
 
     depends_on("mpi", when="+mpi")
-    depends_on("cmake@3.0:", type="build")
+    depends_on("blt")
+    depends_on("cmake")
     depends_on("caliper", when="+caliper")
+    depends_on("chai~examples+raja")
+    depends_on("raja~exercises~examples")
+    depends_on("umpire~examples")
 
     def cmake_args(self):
-        def enabled(variant):
-            return 1 if variant in self.spec else 0
+        spec = self.spec
+        args = []
+
+        args.extend(
+            [
+                "-DCAMP_DIR=%s" % self.spec["camp"].prefix,
+                "-DBLT_SOURCE_DIR=%s" % self.spec["blt"].prefix,
+                "-Dumpire_DIR=%s" % self.spec["umpire"].prefix,
+                "-DRAJA_DIR=%s" % self.spec["raja"].prefix,
+                "-Dchai_DIR=%s" % self.spec["chai"].prefix,
+                "-DENABLE_CHAI=ON",
+            ]
+        )
+
+        if "+caliper" in spec:
+            args.append("-DENABLE_CALIPER=ON")
+
+        if "+mpi" in spec:
+            args.append("-DENABLE_MPI=ON")
+            args.append(self.define("CMAKE_CXX_COMPILER", self.spec["mpi"].mpicxx))
+
+        if "+rocm" in spec:
+            # Set up the hip macros needed by the build
+            args.append("-DENABLE_HIP=ON")
+            args.append("-DHIP_ROOT_DIR={0}".format(spec["hip"].prefix))
+            rocm_archs = spec.variants["amdgpu_target"].value
+            if "none" not in rocm_archs:
+                args.append("-DHIP_HIPCC_FLAGS=--amdgpu-target={0}".format(",".join(rocm_archs)))
+                args.append("-DCMAKE_HIP_ARCHITECTURES={0}".format(rocm_archs))
+        else:
+            # Ensure build with hip is disabled
+            args.append("-DENABLE_HIP=OFF")
+
+        if "+cuda" in spec:
+            args.append("-DENABLE_CUDA=ON")
+            args.append(self.define("CMAKE_CUDA_HOST_COMPILER", self.spec["mpi"].mpicxx))
+            if not spec.satisfies("cuda_arch=none"):
+                cuda_arch = spec.variants["cuda_arch"].value
+                args.append("-DCUDA_ARCH={0}".format(cuda_arch[0]))
+                args.append("-DCMAKE_CUDA_ARCHITECTURES={0}".format(cuda_arch[0]))
+                args.append(
+                    "-DCMAKE_CUDA_FLAGS=--expt-extended-lambda -I%s -I=%s"
+                    % (self.spec["cub"].prefix.include, self.spec["mpi"].prefix.include)
+                )
+        else:
+            args.append("-DENABLE_CUDA=OFF")
 
-        return [
-            "-DENABLE_OPENMP=%d" % enabled("+openmp"),
-            "-DENABLE_MPI=%d" % enabled("+mpi"),
-            "-DENABLE_CALIPER=%d" % enabled("+caliper"),
-        ]
+        return args
 
     def install(self, spec, prefix):
         # Kripke does not provide install target, so we have to copy
         # things into place.
         mkdirp(prefix.bin)
-        install(join_path(self.build_directory, "bin/kripke.exe"), prefix.bin)
+        install(join_path(self.build_directory, "kripke.exe"), prefix.bin)