mvapich2: add pmi_version variant for pmix, pmi1 support (#40665)

var/spack/repos/builtin/packages/mvapich2/package.py

@@ -83,6 +83,15 @@ class Mvapich2(AutotoolsPackage):
         .with_non_feature_values("auto"),
     )
 
+    variant(
+        "pmi_version",
+        description=("The pmi version to be used with slurm"),
+        when="process_managers=slurm",
+        default="pmi2",
+        values=("pmi1", "pmi2", "pmix"),
+        multi=False,
+    )
+
     variant(
         "fabrics",
         description="Select the fabric to be enabled for this build."
@@ -133,6 +142,7 @@ class Mvapich2(AutotoolsPackage):
     depends_on("rdma-core", when="fabrics=nemesisibtcp")
     depends_on("libfabric", when="fabrics=nemesisofi")
     depends_on("slurm", when="process_managers=slurm")
+    depends_on("pmix", when="pmi_version=pmix")
 
     # Fix segmentation fault in `MPIR_Attr_delete_list`:
     # <https://lists.osu.edu/pipermail/mvapich-discuss/2023-January/010695.html>.
@@ -277,11 +287,14 @@ def process_manager_options(self):
 
         # See: http://slurm.schedmd.com/mpi_guide.html#mvapich2
         if "process_managers=slurm" in spec:
-            opts = [
+            opts = ["--with-pm=slurm", "--with-slurm={0}".format(spec["slurm"].prefix)]
-                "--with-pmi=pmi2",
+            if "pmi_version=pmi1" in spec:
-                "--with-pm=slurm",
+                opts.append("--with-pmi=pmi1")
-                "--with-slurm={0}".format(spec["slurm"].prefix),
+            elif "pmi_version=pmi2" in spec:
-            ]
+                opts.append("--with-pmi=pmi2")
+            elif "pmi_version=pmix" in spec:
+                opts.append("--with-pmi=pmix")
+                opts.append("--with-pmix={0}".format(spec["pmix"].prefix))
 
         return opts
 
@@ -347,7 +360,12 @@ def setup_build_environment(self, env):
 
     def setup_run_environment(self, env):
         if "process_managers=slurm" in self.spec:
-            env.set("SLURM_MPI_TYPE", "pmi2")
+            if "pmi_version=pmi1" in self.spec:
+                env.set("SLURM_MPI_TYPE", "pmi1")
+            elif "pmi_version=pmi2" in self.spec:
+                env.set("SLURM_MPI_TYPE", "pmi2")
+            elif "pmi_version=pmix" in self.spec:
+                env.set("SLURM_MPI_TYPE", "pmix")
 
         env.set("MPI_ROOT", self.prefix)