From f8a0d67b2e579cc2bca9c968ebf3f26b8432f1ca Mon Sep 17 00:00:00 2001
From: Markus Kowalewski <markus.kowalewski@gmail.com>
Date: Wed, 17 Jan 2024 16:00:02 +0100
Subject: [PATCH 1/3] mpich: put pmix in build inputs.

The "--with-pmix" option does not require the .dev output.
---
 pkgs/development/libraries/mpich/default.nix | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/libraries/mpich/default.nix b/pkgs/development/libraries/mpich/default.nix
index f53a3268b739..d3bc279506f9 100644
--- a/pkgs/development/libraries/mpich/default.nix
+++ b/pkgs/development/libraries/mpich/default.nix
@@ -36,14 +36,16 @@ stdenv.mkDerivation  rec {
     "FFLAGS=-fallow-argument-mismatch" # https://github.com/pmodels/mpich/issues/4300
     "FCFLAGS=-fallow-argument-mismatch"
   ] ++ lib.optionals pmixSupport [
-    "--with-pmix=${lib.getDev pmix}"
+    "--with-pmix"
   ];
 
   enableParallelBuilding = true;
 
   nativeBuildInputs = [ gfortran python3 ];
   buildInputs = [ perl openssh hwloc ]
-    ++ lib.optional (!stdenv.isDarwin) ch4backend;
+    ++ lib.optional (!stdenv.isDarwin) ch4backend
+    ++ lib.optional pmixSupport pmix;
+
 
   doCheck = true;
 

From a20f3db99526f44703fdf22518db0a524f5e0aea Mon Sep 17 00:00:00 2001
From: Markus Kowalewski <markus.kowalewski@gmail.com>
Date: Wed, 17 Jan 2024 16:25:49 +0100
Subject: [PATCH 2/3] slurm: fix to be able to work with pmix split outputs

---
 pkgs/servers/computing/slurm/default.nix | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/pkgs/servers/computing/slurm/default.nix b/pkgs/servers/computing/slurm/default.nix
index 6e2b2e62d3b6..08fdbcebb016 100644
--- a/pkgs/servers/computing/slurm/default.nix
+++ b/pkgs/servers/computing/slurm/default.nix
@@ -37,6 +37,13 @@ stdenv.mkDerivation rec {
   prePatch = ''
     substituteInPlace src/common/env.c \
         --replace "/bin/echo" "${coreutils}/bin/echo"
+
+    # Autoconf does not support split packages for pmix (libs and headers).
+    # Fix the path to the pmix libraries, so dlopen can find it.
+    substituteInPlace src/plugins/mpi/pmix/mpi_pmix.c \
+        --replace 'xstrfmtcat(full_path, "%s/", PMIXP_LIBPATH)' \
+                  'xstrfmtcat(full_path, "${lib.getLib pmix}/lib/")'
+
   '' + (lib.optionalString enableX11 ''
     substituteInPlace src/common/x11_util.c \
         --replace '"/usr/bin/xauth"' '"${xorg.xauth}/bin/xauth"'
@@ -68,7 +75,7 @@ stdenv.mkDerivation rec {
       "--with-yaml=${lib.getDev libyaml}"
       "--with-ofed=${lib.getDev rdma-core}"
       "--sysconfdir=/etc/slurm"
-      "--with-pmix=${pmix}"
+      "--with-pmix=${lib.getDev pmix}"
       "--with-bpf=${libbpf}"
       "--without-rpath" # Required for configure to pick up the right dlopen path
     ] ++ (optional enableGtk2  "--disable-gtktest")

From 67646d04331ed41831a5c7dcc06fb823a00cc4ce Mon Sep 17 00:00:00 2001
From: Markus Kowalewski <markus.kowalewski@gmail.com>
Date: Wed, 17 Jan 2024 16:35:51 +0100
Subject: [PATCH 3/3] pmix: split output -> out, dev

This reduces the run time closure from 334 MB to 37.8 MB,
mainly by removing gcc from the closure:

nix path-info -hS
/nix/store/jddsvgvf0ayj5x1df9yaa74lhaz6sjz2-pmix-5.0.1     4.1M  334.2M
/nix/store/n6p56psvwpj419f6cjbdlyiri1mv498f-pmix-5.0.1     2.7M   37.8M
---
 pkgs/development/libraries/pmix/default.nix | 29 ++++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/pkgs/development/libraries/pmix/default.nix b/pkgs/development/libraries/pmix/default.nix
index f5fb65f7989b..78605179f710 100644
--- a/pkgs/development/libraries/pmix/default.nix
+++ b/pkgs/development/libraries/pmix/default.nix
@@ -1,5 +1,7 @@
 { lib, stdenv, fetchFromGitHub, perl, autoconf, automake
-, libtool, python3, flex, libevent, hwloc, munge, zlib, pandoc, gitMinimal
+, removeReferencesTo, libtool, python3, flex, libevent
+, targetPackages, makeWrapper
+, hwloc, munge, zlib, pandoc, gitMinimal
 } :
 
 stdenv.mkDerivation rec {
@@ -14,6 +16,8 @@ stdenv.mkDerivation rec {
     fetchSubmodules = true;
   };
 
+  outputs = [ "out" "dev" ];
+
   postPatch = ''
     patchShebangs ./autogen.pl
     patchShebangs ./config
@@ -28,6 +32,8 @@ stdenv.mkDerivation rec {
     flex
     gitMinimal
     python3
+    removeReferencesTo
+    makeWrapper
   ];
 
   buildInputs = [ libevent hwloc munge zlib ];
@@ -46,6 +52,27 @@ stdenv.mkDerivation rec {
 
   postInstall = ''
     find $out/lib/ -name "*.la" -exec rm -f \{} \;
+
+    moveToOutput "bin/pmix_info" "''${!outputDev}"
+    moveToOutput "bin/pmixcc" "''${!outputDev}"
+    moveToOutput "share/pmix/pmixcc-wrapper-data.txt" "''${!outputDev}"
+
+    # The path to the pmixcc-wrapper-data.txt is hard coded and
+    # points to $out instead of dev. Use wrapper to fix paths.
+    wrapProgram $dev/bin/pmixcc \
+      --set PMIX_INCLUDEDIR $dev/include \
+      --set PMIX_PKGDATADIR $dev/share/pmix
+  '';
+
+  postFixup = ''
+    # The build info (parameters to ./configure) are hardcoded
+    # into the library. This clears all references to $dev/include.
+    remove-references-to -t $dev $(readlink -f $out/lib/libpmix.so)
+
+    # Pin the compiler to the current version in a cross compiler friendly way.
+    # Same pattern as for openmpi (see https://github.com/NixOS/nixpkgs/pull/58964#discussion_r275059427).
+    sed -i 's:compiler=.*:compiler=${targetPackages.stdenv.cc}/bin/${targetPackages.stdenv.cc.targetPrefix}cc:' \
+      $dev/share/pmix/pmixcc-wrapper-data.txt
   '';
 
   enableParallelBuilding = true;