From f8dc0f4c1e5625a600d8b238e374115f47df1f1e Mon Sep 17 00:00:00 2001
From: caubet_m <marc.caubet@psi.ch>
Date: Thu, 24 Jul 2025 13:50:48 +0200
Subject: [PATCH] MPICH CPU and GPU variants

---
 Compiler/mpich/build             |  3 +-
 Compiler/mpich/files/config.yaml | 62 ++++++++++++++++++++++++++++----
 Compiler/mpich/modulefile        |  2 +-
 3 files changed, 59 insertions(+), 8 deletions(-)

diff --git a/Compiler/mpich/build b/Compiler/mpich/build
index 8071064..b374342 100755
--- a/Compiler/mpich/build
+++ b/Compiler/mpich/build
@@ -6,9 +6,10 @@ pbuild::pre_configure() {
         pbuild::add_configure_args "--enable-shared"
         pbuild::add_configure_args "--enable-static"
 
-        if [[ -v CUDA_VERSION ]]; then
+        if [[ -v CUDA_VERSION ]] || pbuild::use_flag cuda; then
                 std::info "Enabling CUDA ${CUDA_VERSION}."
                 pbuild::add_configure_args "--with-cuda=${CUDA_HOME}"
+                pbuild::add_configure_args "--with-cuda-lib=${CUDA_HOME}/lib64/stubs"
         fi
 
         if [[ -v HWLOC_VERSION ]]; then
diff --git a/Compiler/mpich/files/config.yaml b/Compiler/mpich/files/config.yaml
index f54042b..d80d933 100644
--- a/Compiler/mpich/files/config.yaml
+++ b/Compiler/mpich/files/config.yaml
@@ -15,19 +15,69 @@ mpich:
   shasums:
     mpich-3.3.2.tar.gz: 4bfaf8837a54771d3e4922c84071ef80ffebddbb6971a006038d91ee7ef959b9
     mpich-4.2.3.tar.gz: 7a019180c51d1738ad9c5d8d452314de65e828ee240bcb2d1f80de9a65be88a8
+    mpich-4.3.1.tar.gz: acc11cb2bdc69678dc8bba747c24a28233c58596f81f03785bf2b7bb7a0ef7dc
 
   versions:
-    4.2.3:
+    4.3.1:
       variants:
-        - systems: [.*.merlin7.psi.ch]
-          relstage: unstable
+        - systems: [login.*.merlin7.psi.ch, cn.*.merlin7.psi.ch]
+          relstage: stable
           overlay: Alps
-          use_overlays: [PSI]
+          target_cpus: ['x86_64'] # MC
+          use_overlays: [PSI, Alps]
           use_flags: [merlin7, slurm]
+          group_deps:
+            compiler: {gcc: [8.5.0,12.3.0]}
+          build_requires: [hwloc/2.12.0, patchelf/0.14.5]
+          runtime_deps: [libfabric/2.2.0-oss]
+          configure_args+:
+            - --with-pm=hydra,gforker
+        - systems: [gpu.*.merlin7.psi.ch]
+          relstage: stable
+          overlay: Alps_A100
+          target_cpus: ['x86_64'] # GPU A100
+          use_overlays: [PSI, Alps_A100]
+          use_flags: [merlin7, slurm, cuda]
+          group_deps:
+            compiler: {gcc: [8.5.0,12.3.0]}
+          build_requires: [hwloc/2.12.0, patchelf/0.14.5]
+          runtime_deps: [cuda/12.9.1,libfabric/2.2.0-oss]
+          configure_args+:
+            - --with-pm=hydra,gforker
+        - systems: [gpu.*.merlin7.psi.ch]
+          relstage: unstable
+          overlay: Alps
+          target_cpus: ['aarch64'] # GPU GH
+          use_overlays: [PSI, Alps]
+          use_flags: [merlin7, slurm, cuda]
           group_deps:
             compiler: {gcc: [12.3.0]}
-          build_requires: [hwloc/2.11.1, libfabric/1.15.2.0, patchelf/0.14.5]
-          runtime_deps: [cuda/12.2.0]
+          build_requires: [hwloc/2.12.0, patchelf/0.14.5]
+          runtime_deps: [cuda/12.9.1,libfabric/2.2.0-oss]
+          configure_args+:
+            - --with-pm=hydra,gforker
+    4.2.3:
+      variants:
+        - systems: ['login.*.merlin7.psi.ch', 'cn.*.merlin7.psi.ch']
+          relstage: stable
+          overlay: Alps
+          use_overlays: [PSI, Alps]
+          use_flags: [merlin7, slurm]
+          group_deps:
+            compiler: {gcc: [8.5.0,12.3.0]}
+          build_requires: [hwloc/2.12.0, patchelf/0.14.5]
+          runtime_deps: [libfabric/2.2.0-oss]
+          configure_args+:
+            - --with-pm=hydra,gforker
+        - systems: [gpu.*.merlin7.psi.ch]
+          relstage: stable
+          overlay: Alps_A100
+          use_overlays: [PSI, Alps_A100]
+          use_flags: [merlin7, slurm, cuda]
+          group_deps:
+            compiler: {gcc: [8.5.0,12.3.0]}
+          build_requires: [hwloc/2.12.0, patchelf/0.14.5]
+          runtime_deps: [cuda/12.9.1,libfabric/2.2.0-oss]
           configure_args+:
             - --with-pm=hydra,gforker
     3.3.2:
diff --git a/Compiler/mpich/modulefile b/Compiler/mpich/modulefile
index 01ede4a..19f27b2 100644
--- a/Compiler/mpich/modulefile
+++ b/Compiler/mpich/modulefile
@@ -3,7 +3,7 @@
 module-whatis		"a high performance implementation of the MPI standard."
 module-url		"http://www.mpich.org/"
 module-license		"BSD-like"
-module-maintainer	"Achim Gsell <achim.gsell@psi.ch>"
+module-maintainer	"Marc Caubet Serrabou <marc.caubet@psi.ch"
 module-help		"
 MPICH is a high-performance and widely portable implementation of the
 MPI Standard, designed to implement all of MPI-1, MPI-2, and MPI-3 (in-