amrex: expand CUDA support (#20740)

* amrex: expand CUDA support

1 files changed, 31 insertions, 4 deletions

diff --git a/var/spack/repos/builtin/packages/amrex/package.py b/var/spack/repos/builtin/packages/amrex/package.py
index 39bf65a081..c0f75af44a 100644
--- a/var/spack/repos/builtin/packages/amrex/package.py
+++ b/var/spack/repos/builtin/packages/amrex/package.py
@@ -6,7 +6,7 @@
 from spack import *
 
 
-class Amrex(CMakePackage):
+class Amrex(CMakePackage, CudaPackage):
     """AMReX is a publicly available software framework designed
     for building massively parallel block- structured adaptive
     mesh refinement (AMR) applications."""
@@ -70,8 +70,6 @@ class Amrex(CMakePackage):
             description='Enable Hypre interfaces')
     variant('petsc', default=False,
             description='Enable PETSc interfaces')
-    variant('cuda', default=False,
-            description='Enable CUDA interfaces')
 
     # Build dependencies
     depends_on('mpi', when='+mpi')
@@ -105,6 +103,16 @@ class Amrex(CMakePackage):
               msg='AMReX PETSc support needs AMReX Fortran API (+fortran)')
     conflicts('+petsc', when='~linear_solvers',
               msg='AMReX PETSc support needs variant +linear_solvers')
+    conflicts('+cuda', when='@:19.08',
+              msg='AMReX CUDA support needs AMReX newer than version 19.08')
+    conflicts('cuda_arch=10', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=11', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=12', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=13', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=20', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=21', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=30', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
+    conflicts('cuda_arch=32', when='+cuda', msg='AMReX only supports compute capabilities >= 3.5')
 
     def url_for_version(self, version):
         if version >= Version('20.05'):
@@ -139,9 +147,17 @@ class Amrex(CMakePackage):
             self.define_from_variant('ENABLE_PETSC', 'petsc'),
             self.define_from_variant('ENABLE_CUDA', 'cuda'),
         ]
+
         if self.spec.satisfies('%fj'):
             args.append('-DCMAKE_Fortran_MODDIR_FLAG=-M')
 
+        if '+cuda' in self.spec:
+            cuda_arch = spec.variants['cuda_arch'].value
+            if cuda_arch == 'none':
+                args.append('-DCUDA_ARCH=Auto')
+            else:
+                args.append('-DCUDA_ARCH={0}'.format(cuda_arch[0]))
+
         return args
 
     #
@@ -167,9 +183,20 @@ class Amrex(CMakePackage):
             self.define_from_variant('AMReX_HDF5', 'hdf5'),
             self.define_from_variant('AMReX_HYPRE', 'hypre'),
             self.define_from_variant('AMReX_PETSC', 'petsc'),
-            self.define_from_variant('AMReX_CUDA', 'cuda'),
         ]
+
         if self.spec.satisfies('%fj'):
             args.append('-DCMAKE_Fortran_MODDIR_FLAG=-M')
 
+        if '+cuda' in self.spec:
+            args.append('-DAMReX_GPU_BACKEND=CUDA')
+            args.append('-DAMReX_CUDA_ERROR_CAPTURE_THIS=ON')
+            args.append('-DAMReX_CUDA_ERROR_CROSS_EXECUTION_SPACE_CALL=ON')
+
+            cuda_arch = spec.variants['cuda_arch'].value
+            if cuda_arch == 'none':
+                args.append('-DAMReX_CUDA_ARCH=Auto')
+            else:
+                args.append('-DAMReX_CUDA_ARCH={0}'.format(cuda_arch[0]))
+
         return args