From e88fc38cfeaf643fce2b9b714921caffa372f615 Mon Sep 17 00:00:00 2001
From: eugeneswalker <38933153+eugeneswalker@users.noreply.github.com>
Date: Fri, 26 Feb 2021 01:17:35 -0800
Subject: [PATCH] hypre: enable cuda variants (#21885)

---
 .../repos/builtin/packages/hypre/package.py   | 41 ++++++++++++++++++-
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/var/spack/repos/builtin/packages/hypre/package.py b/var/spack/repos/builtin/packages/hypre/package.py
index 3f8a892395..61ce3a46dc 100644
--- a/var/spack/repos/builtin/packages/hypre/package.py
+++ b/var/spack/repos/builtin/packages/hypre/package.py
@@ -8,7 +8,7 @@
 import sys
 
 
-class Hypre(Package):
+class Hypre(Package, CudaPackage):
     """Hypre is a library of high performance preconditioners that
        features parallel multigrid methods for both structured and
        unstructured grid problems."""
@@ -78,6 +78,8 @@ class Hypre(Package):
     depends_on("lapack")
     depends_on('superlu-dist', when='+superlu-dist+mpi')
 
+    conflicts('+cuda', when='+int64')
+
     # Patch to build shared libraries on Darwin does not apply to
     # versions before 2.13.0
     conflicts("+shared@:2.12.99 platform=darwin")
@@ -100,7 +102,8 @@ def url_for_version(self, version):
 
         return url.format(version)
 
-    def install(self, spec, prefix):
+    def _configure_args(self):
+        spec = self.spec
         # Note: --with-(lapack|blas)_libs= needs space separated list of names
         lapack = spec['lapack'].libs
         blas = spec['blas'].libs
@@ -162,6 +165,40 @@ def install(self, spec, prefix):
         else:
             configure_args.append("--disable-debug")
 
+        if '+cuda' in self.spec:
+            configure_args.extend([
+                '--with-cuda',
+                '--enable-curand',
+                '--enable-cub'
+            ])
+        else:
+            configure_args.extend([
+                '--without-cuda',
+                '--disable-curand',
+                '--disable-cub'
+            ])
+
+        return configure_args
+
+    def setup_build_environment(self, env):
+        if '+mpi' in self.spec:
+            env.set('CC', self.spec['mpi'].mpicc)
+            env.set('CXX', self.spec['mpi'].mpicxx)
+            env.set('F77', self.spec['mpi'].mpif77)
+
+        if '+cuda' in self.spec:
+            env.set('CUDA_HOME', self.spec['cuda'].prefix)
+            env.set('CUDA_PATH', self.spec['cuda'].prefix)
+            cuda_arch = self.spec.variants['cuda_arch'].value
+            if cuda_arch:
+                arch_sorted = list(sorted(cuda_arch, reverse=True))
+                env.set('HYPRE_CUDA_SM', arch_sorted[0])
+            # In CUDA builds hypre currently doesn't handle flags correctly
+            env.append_flags(
+                'CXXFLAGS', '-O2' if '~debug' in self.spec else '-g')
+
+    def install(self, spec, prefix):
+        configure_args = self._configure_args()
         # Hypre's source is staged under ./src so we'll have to manually
         # cd into it.
         with working_dir("src"):