From 9bae257774e5d97facf9f1b5128069000eb6515f Mon Sep 17 00:00:00 2001
From: Rodrigo Arias Mallo <rodrigo.arias@bsc.es>
Date: Thu, 30 Oct 2025 10:44:27 +0100
Subject: [PATCH] Configure MPICH for MN5

---
 isabel/slurm/flake.nix | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/isabel/slurm/flake.nix b/isabel/slurm/flake.nix
index 9dbb677..d0232d6 100644
--- a/isabel/slurm/flake.nix
+++ b/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