rmgdft: add v5.3.1, v5.4.0 and cuda variant (#37813)

Co-authored-by: Massimiliano Culpo <massimiliano.culpo@gmail.com>

1 files changed, 32 insertions, 11 deletions

diff --git a/var/spack/repos/builtin/packages/rmgdft/package.py b/var/spack/repos/builtin/packages/rmgdft/package.py
index 7f7af92286..dea1afcfe6 100644
--- a/var/spack/repos/builtin/packages/rmgdft/package.py
+++ b/var/spack/repos/builtin/packages/rmgdft/package.py
@@ -7,7 +7,7 @@
 from spack.package import *
 
 
-class Rmgdft(CMakePackage):
+class Rmgdft(CMakePackage, CudaPackage):
     """RMGDFT is a high performance real-space density functional code
     designed for large scale electronic structure calculations."""
 
@@ -16,14 +16,12 @@ class Rmgdft(CMakePackage):
     maintainers("elbriggs")
     tags = ["ecp", "ecp-apps"]
     version("master", branch="master")
+    version("5.4.0", tag="v5.4.0")
+    version("5.3.1", tag="v5.3.1")
     version("5.2.0", tag="v5.2.0")
     version("5.0.5", tag="v5.0.5")
     version("5.0.4", tag="v5.0.4")
     version("5.0.1", tag="v5.0.1")
-    version("4.3.1", tag="v4.3.1")
-    version("4.3.0", tag="v4.3.0")
-    version("4.2.2", tag="v4.2.2")
-    version("4.2.1", tag="v4.2.1")
 
     variant(
         "build_type",
@@ -36,6 +34,8 @@ class Rmgdft(CMakePackage):
 
     variant("local_orbitals", default=True, description="Build O(N) variant.")
 
+    variant("rocm", default=False, description="Build rocm enabled variant.")
+
     # Normally we want this but some compilers (e.g. IBM) are
     # very slow when this is on so provide the option to disable
     variant(
@@ -63,6 +63,11 @@ class Rmgdft(CMakePackage):
     depends_on("fftw-api@3")
     depends_on("mpi")
     depends_on("hdf5")
+    depends_on("cuda", when="+cuda")
+    with when("+rocm"):
+        depends_on("hipblas")
+        depends_on("rocfft")
+        depends_on("rocsolver")
 
     # RMG is a hybrid MPI/threads code and performance is
     # highly dependent on the threading model of the blas
@@ -75,9 +80,18 @@ class Rmgdft(CMakePackage):
     @property
     def build_targets(self):
         spec = self.spec
-        targets = ["rmg-cpu"]
-        if "+local_orbitals" in spec:
-            targets.append("rmg-on-cpu")
+        if "+cuda" in spec:
+            targets = ["rmg-gpu"]
+            cuda_arch_list = spec.variants["cuda_arch"].value
+            cuda_arch = cuda_arch_list[0]
+            if cuda_arch != "none":
+                args.append("-DCUDA_FLAGS=-arch=sm_{0}".format(cuda_arch))
+            if "+local_orbitals" in spec:
+                targets.append("rmg-on-gpu")
+        else:
+            targets = ["rmg-cpu"]
+            if "+local_orbitals" in spec:
+                targets.append("rmg-on-cpu")
         return targets
 
     def cmake_args(self):
@@ -91,6 +105,8 @@ class Rmgdft(CMakePackage):
             args.append("-DUSE_INTERNAL_PSEUDOPOTENTIALS=1")
         else:
             args.append("-DUSE_INTERNAL_PSEUDOPOTENTIALS=0")
+        if "+cuda" in spec:
+            args.append("-DRMG_CUDA_ENABLED=1")
         return args
 
     def install(self, spec, prefix):
@@ -99,9 +115,14 @@ class Rmgdft(CMakePackage):
         mkdirp(prefix.share.tests.RMG)
 
         with working_dir(self.build_directory):
-            install("rmg-cpu", prefix.bin)
-            if "+local_orbitals" in spec:
-                install("rmg-on-cpu", prefix.bin)
+            if "+cuda" in spec:
+                install("rmg-gpu", prefix.bin)
+                if "+local_orbitals" in spec:
+                    install("rmg-on-gpu", prefix.bin)
+            else:
+                install("rmg-cpu", prefix.bin)
+                if "+local_orbitals" in spec:
+                    install("rmg-on-cpu", prefix.bin)
 
         # install tests
         with working_dir(self.build_directory):