Add UCC support

recipe/build-mpi.sh

@@ -33,15 +33,20 @@ if [[ "$target_platform" == osx-* ]]; then
     wrapper_ldflags='-Wl,-rpath,${libdir}'
 fi
 
-# UCX support
+# UCX and UCC support
 build_with_ucx=""
-if [[ "$target_platform" == linux-* ]]; then
+build_with_ucc=""
+if [[ "$target_platform" == linux-* && "$target_platform" != linux-ppc64le ]]; then
+    echo "Build with UCX/UCC support"
     build_with_ucx="--with-ucx=$PREFIX"
+    build_with_ucc="--with-ucc=$PREFIX"
 fi
 
+
 # CUDA support
 build_with_cuda=""
 if [[ -n "$CUDA_HOME" ]]; then
+    echo "Build with CUDA support"
     build_with_cuda="--with-cuda=$CUDA_HOME --with-cuda-libdir=$CUDA_HOME/lib64/stubs"
 fi
 
@@ -81,6 +86,7 @@ fi
             --enable-mca-dso \
             --enable-ipv6 \
             $build_with_ucx \
+            $build_with_ucc \
             $build_with_cuda \
     || (cat config.log; false)
 
@@ -95,6 +101,11 @@ if [ -n "$build_with_ucx" ]; then
     echo "osc = ^ucx" >> $PREFIX/etc/openmpi-mca-params.conf
     cat $RECIPE_DIR/post-link-ucx.sh >> $POST_LINK
 fi
+if [ -n "$build_with_ucc" ]; then
+    echo "setting MCA coll_ucc_enable to 0..."
+    echo "coll_ucc_enable = 0" >> $PREFIX/etc/openmpi-mca-params.conf
+    cat $RECIPE_DIR/post-link-ucc.sh >> $POST_LINK
+fi
 if [ -n "$build_with_cuda" ]; then
     echo "setting MCA mca_base_component_show_load_errors to 0..."
     echo "mca_base_component_show_load_errors = 0" >> $PREFIX/etc/openmpi-mca-params.conf

recipe/meta.yaml

@@ -1,7 +1,7 @@
 {% set version = "5.0.5" %}
 {% set major = version.rpartition('.')[0] %}
 {% set cuda_major = (cuda_compiler_version|default("11.8")).rpartition('.')[0] %}
-{% set build = 0 %}
+{% set build = 1 %}
 
 # give conda package a higher build number
 {% if mpi_type == 'conda' %}
@@ -41,6 +41,7 @@ outputs:
       {% endif %}
       ignore_run_exports:
         - ucx  # [linux]
+        - ucc  # [linux]
       ignore_run_exports_from:
         - {{ compiler('cuda') }}   # [cuda_compiler != "None"]
     requirements:
@@ -54,23 +55,25 @@ outputs:
         - perl  # [osx]
         - autoconf  # [osx]
         - automake  # [osx]
-        - libtool   # [unix]
+        - libtool  # [unix]
         - make  # [unix]
       host:
         #- openpmix
         #- prrte
-        - libhwloc
         - libevent
+        - libhwloc
         - libnl  # [linux]
         - zlib
-        - ucx  # [linux]
+        - ucc  # [linux and not ppc64le]
+        - ucx  # [linux and not ppc64le]
         - cuda-version {{ cuda_compiler_version }}  # [cuda_compiler != "None"]
       run:
         - mpi 1.0 openmpi
         #- openpmix
         #- prrte
       run_constrained:
-        - {{ pin_compatible("ucx", max_pin="x.x") }}  # [linux]
+        - {{ pin_compatible("ucx", max_pin="x.x") }}  # [linux and not ppc64le]
+        - {{ pin_compatible("ucc", max_pin="x.x") }}  # [linux and not ppc64le]
         # Open MPI only uses CUDA Driver APIs, set the minimal driver version
         - __cuda  >= {{ cuda_major ~ ".0" }}  # [cuda_compiler != "None"]
         # Ensure a consistent CUDA environment
@@ -158,3 +161,4 @@ extra:
     - msarahan
     - ocefpaf
     - beckermr
+    - j34ni

recipe/post-link-ucc.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+cat << EOF >> $PREFIX/.messages.txt
+
+On Linux, Open MPI is built with UCC support but it is disabled by default.
+To enable it, first install UCC (conda install -c conda-forge ucc).
+Afterwards, set the environment variables
+OMPI_MCA_coll_ucc_enable=1
+before launching your MPI processes.
+Equivalently, you can set the MCA parameters in the command line:
+mpiexec --mca coll_ucc_enable 1 ...
+
+EOF

recipe/run_test.sh

@@ -12,11 +12,25 @@ if [[ $PKG_NAME == "openmpi" ]]; then
     exit 1
   fi
 
+  if [[ "$target_platform" == linux-64 || "$target_platform" == linux-aarch64 ]]; then
+     if [[ -z "$(ompi_info | grep ucx)" ]]; then
+        echo "OpenMPI configured without UCX support!"
+        exit 1
+     fi
+  fi
+
   if [[ -n "$(conda list | grep cuda-version)" ]]; then
     echo "Improper CUDA dependency!"
     exit 1
   fi
 
+  if [[ "$target_platform" == linux-64 || "$target_platform" == linux-aarch64 ]]; then
+     if [[ -z "$(ompi_info | grep cuda)" ]]; then
+        echo "OpenMPI configured without CUDA support!"
+        exit 1
+     fi
+  fi
+
   command -v ompi_info
   ompi_info