enhance psmpi easyblock to activate CUDA support when it is a dependency + make check for static libs in MPICH easyblock optional

easybuild/easyblocks/m/mpich.py

@@ -134,7 +134,7 @@ def configure_step(self, add_mpich_configopts=True):
 
     # make and make install are default
 
-    def sanity_check_step(self, custom_paths=None, use_new_libnames=None, check_launchers=True):
+    def sanity_check_step(self, custom_paths=None, use_new_libnames=None, check_launchers=True, check_static_libs=True):
         """
         Custom sanity check for MPICH
         """
@@ -160,7 +160,10 @@ def sanity_check_step(self, custom_paths=None, use_new_libnames=None, check_laun
 
         bins = [os.path.join('bin', x) for x in binaries]
         headers = [os.path.join('include', x) for x in ['mpi.h', 'mpicxx.h', 'mpif.h']]
-        libs_fn = ['lib%s.%s' % (libname, e) for libname in libnames for e in ['a', shlib_ext]]
+        lib_exts = [shlib_ext]
+        if check_static_libs:
+            lib_exts.append('a')
+        libs_fn = ['lib%s.%s' % (lib, e) for lib in libnames for e in lib_exts]
         libs = [(os.path.join('lib', lib), os.path.join('lib64', lib)) for lib in libs_fn]
 
         custom_paths.setdefault('dirs', []).extend(['bin', 'include', ('lib', 'lib64')])

easybuild/easyblocks/p/psmpi.py

@@ -28,6 +28,7 @@
 @author: Damian Alvarez (Forschungszentrum Juelich)
 """
 
+import easybuild.tools.environment as env
 import easybuild.tools.toolchain as toolchain
 
 from distutils.version import LooseVersion
@@ -58,6 +59,7 @@ def extra_options(extra_vars=None):
             'mpich_opts': [None, "Optional options to configure MPICH", CUSTOM],
             'threaded': [False, "Enable multithreaded build (which is slower)", CUSTOM],
             'pscom_allin_path': [None, "Enable pscom integration by giving its source path", CUSTOM],
+            'cuda': [False, "Enable CUDA awareness", CUSTOM],
         })
         return extra_vars
 
@@ -72,10 +74,20 @@ def configure_step(self):
         comp_opts = {
             toolchain.GCC: 'gcc',
             toolchain.INTELCOMP: 'intel',
-            # TODO: Include PGI as soon as it is available as toolchain
-            # toolchain.PGI: 'pgi',
+            toolchain.PGI: 'pgi',
+            toolchain.NVHPC: 'nvhpc',
         }
 
+        # ParaStationMPI defines its environment through confsets. So these should be unset
+        env_vars = ['CFLAGS', 'CPPFLAGS', 'CXXFLAGS', 'FCFLAGS', 'FFLAGS', 'LDFLAGS', 'LIBS']
+        env.unset_env_vars(env_vars)
+        self.log.info("Unsetting the following variables: " + ' '.join(env_vars))
+
+        # Enable CUDA
+        if self.cfg['cuda']:
+            self.log.info("Enabling CUDA-Awareness...")
+            self.cfg.update('configopts', ' --with-cuda')
+
         # Set confset
         comp_fam = self.toolchain.comp_family()
         if comp_fam in comp_opts:
@@ -99,7 +111,8 @@ def configure_step(self):
             pscom_path = self.cfg['pscom_allin_path'].strip()
             self.cfg.update('configopts', ' --with-pscom-allin="%s"' % pscom_path)
 
-        pscom_flags = 'PSCOM_LDFLAGS=-L{0}/lib PSCOM_CPPFLAGS=-I{0}/include'.format(pscom_path)
+        pscom_flags = 'export PSCOM_LDFLAGS="-L{0}/lib $PSCOM_LDFLAGS" &&'.format(pscom_path)
+        pscom_flags += ' export PSCOM_CPPFLAGS="-I{0}/include $PSCOM_CPPFLAGS" &&'.format(pscom_path)
         self.cfg.update('preconfigopts', pscom_flags)
 
         super(EB_psmpi, self).configure_step(add_mpich_configopts=False)
@@ -117,4 +130,6 @@ def sanity_check_step(self):
         # ParaStationMPI < 5.1.1-1 is based on MPICH < 3.1.1.
         use_new_libnames = LooseVersion(self.version) >= LooseVersion('5.1.1-1')
 
-        super(EB_psmpi, self).sanity_check_step(use_new_libnames=use_new_libnames, check_launchers=False)
+        super(EB_psmpi, self).sanity_check_step(use_new_libnames=use_new_libnames,
+                                                check_launchers=False,
+                                                check_static_libs=False)