intel-oneapi-compilers/mpi: add module support (#20808)

Facilitate running intel-oneapi-mpi outside of Spack (set PATH,
LD_LIBRARY_PATH, etc. appropriately).

Co-authored-by: Robert Cohn <rscohn2@gmail.com>

2 files changed, 48 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/intel-oneapi-compilers/package.py b/var/spack/repos/builtin/packages/intel-oneapi-compilers/package.py
index 3ad097976f..6a90e10621 100644
--- a/var/spack/repos/builtin/packages/intel-oneapi-compilers/package.py
+++ b/var/spack/repos/builtin/packages/intel-oneapi-compilers/package.py
@@ -79,3 +79,8 @@ class IntelOneapiCompilers(IntelOneApiPackage):
         env.set('CC', self._join_prefix('bin/icx'))
         env.set('CXX', self._join_prefix('bin/icpx'))
         env.set('FC', self._join_prefix('bin/ifx'))
+        # Set these so that MPI wrappers will pick up these compilers
+        # when this module is loaded.
+        env.set('I_MPI_CC', 'icx')
+        env.set('I_MPI_CXX', 'icpx')
+        env.set('I_MPI_FC', 'ifx')
diff --git a/var/spack/repos/builtin/packages/intel-oneapi-mpi/package.py b/var/spack/repos/builtin/packages/intel-oneapi-mpi/package.py
index 5e22d0de29..1f7aef1403 100644
--- a/var/spack/repos/builtin/packages/intel-oneapi-mpi/package.py
+++ b/var/spack/repos/builtin/packages/intel-oneapi-mpi/package.py
@@ -4,6 +4,8 @@
 # SPDX-License-Identifier: (Apache-2.0 OR MIT)
 
 
+import subprocess
+
 from spack import *
 
 releases = {
@@ -21,6 +23,8 @@ class IntelOneapiMpi(IntelOneApiLibraryPackage):
 
     provides('mpi@:3')
 
+    depends_on('patchelf', type='build')
+
     def __init__(self, spec):
         self.component_info(dir_name='mpi',
                             components='intel.oneapi.lin.mpi.devel',
@@ -50,3 +54,42 @@ class IntelOneapiMpi(IntelOneApiLibraryPackage):
             ldir = find_libraries('*', root=lib_path, shared=True, recursive=False)
             libs += ldir
         return libs
+
+    def _join_prefix(self, path):
+        return join_path(self.prefix, 'mpi', 'latest', path)
+
+    def _ld_library_path(self):
+        dirs = ['lib',
+                'lib/release',
+                'libfabric/lib']
+        for dir in dirs:
+            yield self._join_prefix(dir)
+
+    def _library_path(self):
+        dirs = ['lib',
+                'lib/release',
+                'libfabric/lib']
+        for dir in dirs:
+            yield self._join_prefix(dir)
+
+    def install(self, spec, prefix):
+        super(IntelOneapiMpi, self).install(spec, prefix)
+
+        # need to patch libmpi.so so it can always find libfabric
+        libfabric_rpath = self._join_prefix('libfabric/lib')
+        for lib_version in ['debug', 'release', 'release_mt', 'debug_mt']:
+            file = self._join_prefix('lib/' + lib_version + '/libmpi.so')
+            subprocess.call(['patchelf', '--set-rpath', libfabric_rpath, file])
+
+    def setup_run_environment(self, env):
+        env.prepend_path('PATH', self._join_prefix('bin'))
+        env.prepend_path('CPATH', self._join_prefix('include'))
+        for dir in self._library_path():
+            env.prepend_path('LIBRARY_PATH', dir)
+        for dir in self._ld_library_path():
+            env.prepend_path('LD_LIBRARY_PATH', dir)
+        # so wrappers know where MPI lives
+        mpi_root = join_path(prefix, 'mpi', 'latest')
+        env.set('I_MPI_ROOT', mpi_root)
+        # set this so that wrappers can find libfabric providers
+        env.set('FI_PROVIDER_PATH', self._join_prefix('libfabric/lib/prov'))