mpich: add support for Mellanox HCOLL (#30662)

Co-authored-by: Federico Ficarelli <federico.ficarelli@pm.me>

1 files changed, 12 insertions, 0 deletions

diff --git a/var/spack/repos/builtin/packages/mpich/package.py b/var/spack/repos/builtin/packages/mpich/package.py
index b758f5b720..64fd90f0d6 100644
--- a/var/spack/repos/builtin/packages/mpich/package.py
+++ b/var/spack/repos/builtin/packages/mpich/package.py
@@ -115,6 +115,12 @@ with '-Wl,-commons,use_dylibs' and without
     depends_on('yaksa+rocm', when='+rocm ^yaksa')
     conflicts('datatype-engine=yaksa', when='device=ch3')
 
+    variant('hcoll', default=False,
+            description='Enable support for Mellanox HCOLL accelerated '
+                        'collective operations library',
+            when='@3.3: device=ch4 netmod=ucx')
+    depends_on('hcoll', when='+hcoll')
+
     # Todo: cuda can be a conditional variant, but it does not seem to work when
     # overriding the variant from CudaPackage.
     conflicts('+cuda', when='@:3.3')
@@ -345,6 +351,9 @@ with '-Wl,-commons,use_dylibs' and without
             if match:
                 variants.append('netmod=' + match.group(1))
 
+            if re.search(r'--with-hcoll', output):
+                variants += '+hcoll'
+
             match = re.search(r'MPICH CC:\s+(\S+)', output)
             compiler_spec = get_spack_compiler_spec(
                 os.path.dirname(match.group(1)))
@@ -559,6 +568,9 @@ with '-Wl,-commons,use_dylibs' and without
         elif 'datatype-engine=auto' in spec:
             config_args.append('--with-datatye-engine=auto')
 
+        if '+hcoll' in spec:
+            config_args.append('--with-hcoll=' + spec['hcoll'].prefix)
+
         return config_args
 
     @run_after('install')