5 files changed, 85 insertions, 0 deletions
diff --git a/doc/hooks/index.md b/doc/hooks/index.md
index 602febaf9d9ba..8100e91c8b481 100644
--- a/doc/hooks/index.md
+++ b/doc/hooks/index.md
@@ -17,6 +17,7 @@ installShellFiles.section.md
 libiconv.section.md
 libxml2.section.md
 meson.section.md
+mpi-check-hook.section.md
 ninja.section.md
 patch-rc-path-hooks.section.md
 perl.section.md
diff --git a/doc/hooks/mpi-check-hook.section.md b/doc/hooks/mpi-check-hook.section.md
new file mode 100644
index 0000000000000..e3fb5c40dadaa
--- /dev/null
+++ b/doc/hooks/mpi-check-hook.section.md
@@ -0,0 +1,24 @@
+#  mpiCheckPhaseHook {#setup-hook-mpi-check}
+
+
+This hook can be used to setup a check phase that
+requires running a MPI application. It detects the
+used present MPI implementaion type and exports
+the neceesary environment variables to use
+`mpirun` and `mpiexec` in a Nix sandbox.
+
+
+Example:
+
+```nix
+  { mpiCheckPhaseHook, mpi, ... }:
+
+  ...
+
+  nativeCheckInputs = [
+    openssh
+    mpiCheckPhaseHook
+  ];
+```
+
+
diff --git a/pkgs/build-support/setup-hooks/mpi-check-hook/default.nix b/pkgs/build-support/setup-hooks/mpi-check-hook/default.nix
new file mode 100644
index 0000000000000..2834cfcc44ff1
--- /dev/null
+++ b/pkgs/build-support/setup-hooks/mpi-check-hook/default.nix
@@ -0,0 +1,5 @@
+{ callPackage, makeSetupHook }:
+
+makeSetupHook {
+  name = "mpi-checkPhase-hook";
+} ./mpi-check-hook.sh
diff --git a/pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh b/pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh
new file mode 100644
index 0000000000000..fca1f7b7f9327
--- /dev/null
+++ b/pkgs/build-support/setup-hooks/mpi-check-hook/mpi-check-hook.sh
@@ -0,0 +1,54 @@
+preCheckHooks+=('setupMpiCheck')
+preInstallCheckHooks+=('setupMpiCheck')
+
+
+setupMpiCheck() {
+  # Find out which MPI implementation we are using
+  # and set safe defaults that are guaranteed to run
+  # on any build machine
+
+  mpiType="NONE"
+
+  # OpenMPI signature
+  if command ompi_info &> /dev/null; then
+    mpiType="openmpi"
+  fi
+
+  # MPICH based implementations
+  if command mpichversion &> /dev/null; then
+    if [ "$mpiType" != "NONE" ]; then
+      echo "WARNING: found OpenMPI and MPICH/MVAPICH executables"
+    fi
+
+    version=$(mpichversion)
+    if [[ "$version" == *"MPICH"* ]]; then
+      mpiType="MPICH"
+    fi
+    if [[ "$version" == *"MVAPICH"* ]]; then
+      mpiType="MVAPICH"
+    fi
+  fi
+
+  echo "Found MPI implementation: $mpiType"
+
+  case $mpiType in
+    openmpi)
+      # make sure the test starts even if we have less than the requested amount of cores
+      export OMPI_MCA_rmaps_base_oversubscribe=1
+      # Disable CPU pinning
+      export OMPI_MCA_hwloc_base_binding_policy=none
+      ;;
+    MPICH)
+      # Fix to make mpich run in a sandbox
+      export HYDRA_IFACE=lo
+      ;;
+    MVAPICH)
+      # Disable CPU pinning
+      export MV2_ENABLE_AFFINITY=0
+      ;;
+  esac
+
+  # Limit number of OpenMP threads. Default is "all cores".
+  export OMP_NUM_THREADS=1
+}
+
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index c89e688e2db64..17f9d307d59ca 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -12258,6 +12258,7 @@ with pkgs;
   outils = callPackage ../tools/misc/outils { };
 
   mpi = openmpi; # this attribute should used to build MPI applications
+  mpiCheckPhaseHook = callPackage ../build-support/setup-hooks/mpi-check-hook { };
 
   ucc = callPackage ../development/libraries/ucc { };