Configure MPICH for MN5

isabel/slurm/flake.nix

@@ -3,6 +3,7 @@
   outputs = { self, jungle }:
   let
     nixpkgs = jungle.inputs.nixpkgs;
+    lib = nixpkgs.lib;
     clusterOverlay = (final: prev: {
       # Use cluster llvm compiler repo
       clangOmpss2Unwrapped = prev.clangOmpss2Unwrapped.override {
@@ -11,6 +12,13 @@
         gitCommit = "151c260ba834826c01855da0a41fc203ffe4d025";
         gitBranch = "cluster";
       };
+      # Configure MPICH to use UCX with multiple thread support
+      mpich = prev.mpich.overrideAttrs (old: {
+        configureFlags = (lib.remove "--with-device=ch4:ofi" old.configureFlags) ++ [
+          "--with-device=ch4:ucx"
+          "--enable-threads=multiple"
+        ];
+      });
     });
     pkgs = import nixpkgs {
       system = "x86_64-linux";
@@ -26,7 +34,7 @@
       pname = "devshell";
       buildInputs = with pkgs; [
         slurm.out slurm.dev gcc
-        clangOmpss2
+        clangOmpss2 mpich osumb
       ];
       inputsFrom = with pkgs; [
         nanos6