WarpX: MPI, Backends, Ascent (#15647)

Expose serial/parallel build (MPI), CUDA/OpenMP backends, Clang, and
Ascent bindings.

Interestingly, `warpx +ascent` currently leads to an infinite loop in
the Spack concretizer.

1 files changed, 45 insertions, 20 deletions

diff --git a/var/spack/repos/builtin/packages/warpx/package.py b/var/spack/repos/builtin/packages/warpx/package.py
index a729f9ac28..55cdf3cb13 100644
--- a/var/spack/repos/builtin/packages/warpx/package.py
+++ b/var/spack/repos/builtin/packages/warpx/package.py
@@ -24,21 +24,32 @@ class Warpx(MakefilePackage):
 
     variant('dims',
             default='3',
-            values=('2', '3'),
+            values=('2', '3', 'rz'),
             multi=False,
             description='Number of spatial dimensions')
-
+    variant('backend',
+            default='openmp',
+            values=('openmp', 'cuda', 'hip'),
+            multi=True,
+            description='Programming model for compute kernels')
+    variant('mpi', default=True, description='Enable MPI support')
     variant('psatd', default=False, description='Enable PSATD solver')
     variant('debug', default=False, description='Enable debugging features')
     variant('tprof', default=True, description='Enable tiny profiling features')
-    variant('openmp', default=True, description='Enable OpenMP features')
     variant('openpmd', default=True, description='Enable openPMD I/O')
+    variant('ascent', default=False, description='Enable Ascent in situ vis')
 
-    depends_on('mpi')
+    depends_on('cuda', when='backend=cuda')
+    depends_on('mpi', when='+mpi')
     depends_on('fftw@3:', when='+psatd')
+    depends_on('fftw +mpi', when='+psatd +mpi')
     depends_on('pkgconfig', type='build', when='+openpmd')
-    depends_on('python', type='build')  # FIXME upstream
-    depends_on('openpmd-api@0.11.0:,dev +mpi', when='+openpmd')
+    depends_on('python', type='build')  # AMReX' build system info
+    depends_on('openpmd-api@0.11.0:,dev', when='+openpmd')
+    depends_on('openpmd-api +mpi', when='+openpmd +mpi')
+    depends_on('ascent', when='+ascent')
+    depends_on('ascent +cuda', when='+ascent backend=cuda')
+    depends_on('ascent +mpi ^conduit~hdf5', when='+ascent +mpi')
 
     resource(name='amrex',
              git='https://github.com/AMReX-Codes/amrex.git',
@@ -49,36 +60,50 @@ class Warpx(MakefilePackage):
              git='https://bitbucket.org/berkeleylab/picsar.git',
              tag='master')
 
+    conflicts('backend=cuda', when='backend=hip',
+              msg='WarpX can be compiled with either CUDA or HIP backend')
+    conflicts('backend=hip', msg='WarpX\' HIP backend is not yet implemented')
+
     def edit(self, spec, prefix):
         comp = 'gcc'
-        vendors = {'%gcc': 'gcc', '%intel': 'intel'}
+        vendors = {'%gcc': 'gcc', '%intel': 'intel', '%clang': 'llvm'}
         for key, value in vendors.items():
             if self.spec.satisfies(key):
                 comp = value
 
-        def torf(s):
-            "Returns the string TRUE or FALSE"
-            return repr(s in spec).upper()
+        # Returns the string TRUE or FALSE
+        torf = lambda s: repr(s in spec).upper()
+
+        try:
+            self.compiler.openmp_flag
+        except UnsupportedCompilerFlag:
+            use_omp = 'FALSE'
+        else:
+            use_omp = torf('backend=openmp')
 
         makefile = FileFilter('GNUmakefile')
         makefile.filter('AMREX_HOME .*', 'AMREX_HOME = amrex')
         makefile.filter('PICSAR_HOME .*', 'PICSAR_HOME = picsar')
         makefile.filter('COMP .*', 'COMP = {0}'.format(comp))
-        makefile.filter('DIM .*',
-                        'DIM = {0}'.format(int(spec.variants['dims'].value)))
+        makefile.filter('USE_MPI .*',
+                        'USE_MPI = {0}'.format(torf('+mpi')))
+        if 'dims=rz' in spec:
+            makefile.filter('USE_RZ .*', 'USE_RZ = TRUE')
+        else:
+            makefile.filter('DIM .*', 'DIM = {0}'.format(
+                int(spec.variants['dims'].value)))
         makefile.filter('USE_PSATD .*',
                         'USE_PSATD = {0}'.format(torf('+psatd')))
-        try:
-            self.compiler.openmp_flag
-        except UnsupportedCompilerFlag:
-            use_omp = 'FALSE'
-        else:
-            use_omp = torf('+openmp')
-        use_openpmd = torf('+openpmd')
         makefile.filter('USE_OMP .*',
                         'USE_OMP = {0}'.format(use_omp))
+        makefile.filter('USE_GPU .*',
+                        'USE_GPU = {0}'.format(torf('backend=cuda')))
+        makefile.filter('USE_HIP .*',
+                        'USE_HIP = {0}'.format(torf('backend=hip')))
         makefile.filter('USE_OPENPMD .*',
-                        'USE_OPENPMD = {0}'.format(use_openpmd))
+                        'USE_OPENPMD = {0}'.format(torf('+openpmd')))
+        makefile.filter('USE_ASCENT_INSITU .*',
+                        'USE_ASCENT_INSITU = {0}'.format(torf('+ascent')))
         makefile.filter('DEBUG .*',
                         'DEBUG = {0}'.format(torf('+debug')))
         makefile.filter('TINY_PROFILE .*',