diff --git a/.github/workflows/cd-wheel.yml b/.github/workflows/cd-wheel.yml index e16b4e7..6d3e30a 100644 --- a/.github/workflows/cd-wheel.yml +++ b/.github/workflows/cd-wheel.yml @@ -76,6 +76,7 @@ jobs: shell: python name: 'setup build matrix' run: | + # setup build matrix # " keys = ("os", "arch", "runner") rows = [ ("Linux", "aarch64", "ubuntu-22.04"), @@ -93,6 +94,7 @@ jobs: import os, json with open(os.getenv("GITHUB_OUTPUT"), "w") as out: print(f"matrix={json.dumps(matrix)}", file=out) + # " build: needs: setup @@ -117,12 +119,12 @@ jobs: cd $(brew --prefix)/bin gfortran=$(ls gfortran-* | sort | head -n 1) sudo ln -s $gfortran gfortran + # unlink libevent + brew unlink libevent || true # install autotools brew install autoconf brew install automake brew install libtool - # unlink libevent - brew unlink libevent || true # install uv brew install uv @@ -142,9 +144,13 @@ jobs: - id: source-date-epoch run: | - SOURCE_DATE_EPOCH=$(git log -1 --pretty=%ct) + read -r SOURCE_DATE_EPOCH < source-date-epoch || true + SOURCE_DATE_EPOCH=${SOURCE_DATE_EPOCH:-$(git log -1 --pretty=%ct)} echo SOURCE_DATE_EPOCH=$SOURCE_DATE_EPOCH >> $GITHUB_ENV - echo $(git log -1 --pretty=%ci) [timestamp=$SOURCE_DATE_EPOCH] + echo [SOURCE_DATE_EPOCH=$SOURCE_DATE_EPOCH] $( + date -u -d @$SOURCE_DATE_EPOCH 2>/dev/null || + date -u -r $SOURCE_DATE_EPOCH 2>/dev/null ) + working-directory: package/source - id: build uses: pypa/cibuildwheel@v2.21.3 @@ -190,7 +196,9 @@ jobs: - id: upload uses: actions/upload-artifact@v4 with: - name: wheel-${{ inputs.mpiname }}-${{ matrix.os }}-${{ matrix.arch }} + name: "wheel-${{ inputs.mpiname }}-\ + ${{ inputs.version || 'latest' }}-\ + ${{ matrix.os }}-${{ matrix.arch }}" path: wheelhouse/*.whl - id: check diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index aea16d1..9f68586 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -11,12 +11,32 @@ on: # yamllint disable-line rule:truthy jobs: - mpich: + mpich-42: uses: ./.github/workflows/cd-wheel.yml with: mpiname: mpich + version: - openmpi: + mpich-41: + uses: ./.github/workflows/cd-wheel.yml + with: + mpiname: mpich + version: 4.1.3 + + mpich-34: + uses: ./.github/workflows/cd-wheel.yml + with: + mpiname: mpich + version: 3.4.3 + + openmpi-50: + uses: ./.github/workflows/cd-wheel.yml + with: + mpiname: openmpi + version: + + openmpi-41: uses: ./.github/workflows/cd-wheel.yml with: mpiname: openmpi + version: 4.1.6 diff --git a/.gitignore b/.gitignore index f5c168b..f654289 100644 --- a/.gitignore +++ b/.gitignore @@ -5,4 +5,5 @@ /venv/ *.egg-info/ *.tar.gz +*.tar.bz2 *~ diff --git a/Makefile b/Makefile index eadfffa..5e2fe82 100644 --- a/Makefile +++ b/Makefile @@ -10,10 +10,11 @@ lint: yamllint .github/ clean: + $(RM) -r package/METADATA + $(RM) -r package/LICENSE* $(RM) -r package/build - $(RM) -r package/LICENSE - $(RM) -r package/install $(RM) -r package/source $(RM) -r package/workdir + $(RM) -r package/install $(RM) -r package/*.egg-info $(RM) -r .*_cache diff --git a/bootstrap.sh b/bootstrap.sh index 34e918b..e908c59 100755 --- a/bootstrap.sh +++ b/bootstrap.sh @@ -8,6 +8,11 @@ case "$mpiname" in esac version=${VERSION:-$version} +ucxversion=1.17.0 +ofiversion=1.22.0 +ucxversion=${UCXVERSION:-$ucxversion} +ofiversion=${OFIVERSION:-$ofiversion} + PROJECT=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) PACKAGE=$PROJECT/package SOURCE=$PACKAGE/source @@ -27,21 +32,44 @@ fi if test ! -d "$SOURCE"; then if test ! -f "$tarball"; then echo downloading "$urlbase"/"$tarball"... - curl -fsO "$urlbase"/"$tarball" + curl -fsSLO "$urlbase"/"$tarball" else echo reusing "$tarball"... fi - echo extracting "$tarball" to "$SOURCE"... + echo extracting "$tarball"... tar xf "$tarball" mv "$mpiname-$version" "$SOURCE" patch="$PROJECT/patches/$mpiname-$version" if test -f "$patch"; then patch -p1 -i "$patch" -d "$SOURCE" fi - if test "$mpiname-$(uname)" = "openmpi-Darwin"; then - if test -d "$SOURCE"/3rd-party; then - cd "$SOURCE"/3rd-party - tar xf libevent-*.tar.gz && cd libevent-* + if test "$mpiname" = "mpich"; then + if test "${version}" \< "4.2.0"; then + disable_doc='s/^\(install-data-local:\s\+\)\$/\1#\$/' + Makefile="$SOURCE"/Makefile.in + sed -i.orig "$disable_doc" "$Makefile" + fi + fi + if test "$mpiname" = "openmpi"; then + for deptarball in "$SOURCE"/3rd-party/*.tar.*; do + test -f "$deptarball" || continue + echo extracting "$(basename "$deptarball")" + tar xf "$deptarball" -C "$(dirname "$deptarball")" + done + makefiles=( + "$SOURCE"/3rd-party/openpmix/src/util/keyval/Makefile.in + "$SOURCE"/3rd-party/prrte/src/mca/rmaps/rank_file/Makefile.in + "$SOURCE"/3rd-party/prrte/src/util/hostfile/Makefile.in + ) + for makefile in "${makefiles[@]}"; do + test -f "$makefile" || continue + echo "PMIX_CFLAGS_BEFORE_PICKY = @CFLAGS@" >> "$makefile" + echo "PRTE_CFLAGS_BEFORE_PICKY = @CFLAGS@" >> "$makefile" + done + fi + if test "$mpiname" = "openmpi" && test "${version}" \< "5.0.5"; then + if test "$(uname)" = "Darwin" && test -d "$SOURCE"/3rd-party; then + cd "$SOURCE"/3rd-party/libevent-* echo running autogen.sh on "$(basename "$(pwd)")" ./autogen.sh cd "$PROJECT" @@ -49,6 +77,115 @@ if test ! -d "$SOURCE"; then fi else echo reusing directory "$SOURCE"... + check() { test "$(awk "/$1/"'{print $2}' "$PACKAGE/METADATA")" = "$2"; } + check Name "$mpiname" || (echo not "$mpiname-$version"!!! && exit 1) + check Version "$version" || (echo not "$mpiname-$version"!!! && exit 1) +fi + +if test "$(uname)" = "Linux"; then + case "$mpiname" in + mpich) MODSOURCE="$SOURCE"/modules ;; + openmpi) MODSOURCE="$SOURCE"/3rd-party ;; + esac + ofigithub="https://github.com/ofiwg/libfabric" + ofiurlbase="$ofigithub/releases/download/v$ofiversion" + ofitarball="libfabric-$ofiversion.tar.bz2" + ofidestdir="$MODSOURCE"/"${ofitarball%%.tar.*}" + if test ! -d "$ofidestdir"; then + if test ! -f "$ofitarball"; then + echo downloading "$ofiurlbase"/"$ofitarball"... + curl -fsSLO "$ofiurlbase"/"$ofitarball" + else + echo reusing "$ofitarball"... + fi + echo extracting "$ofitarball"... + tar xf "$ofitarball" + mkdir -p "$(dirname "$ofidestdir")" + mv "$(basename "$ofidestdir")" "$ofidestdir" + else + echo reusing directory "$ofidestdir"... + fi + ucxgithub="https://github.com/openucx/ucx" + ucxurlbase="$ucxgithub/releases/download/v$ucxversion" + ucxtarball="ucx-$ucxversion.tar.gz" + ucxdestdir="$MODSOURCE"/"${ucxtarball%%.tar.*}" + if test ! -d "$ucxdestdir"; then + if test ! -f "$ucxtarball"; then + echo downloading "$ucxurlbase"/"$ucxtarball"... + curl -fsSLO "$ucxurlbase"/"$ucxtarball" + else + echo reusing "$ucxtarball"... + fi + echo extracting "$ucxtarball"... + tar xf "$ucxtarball" + mkdir -p "$(dirname "$ucxdestdir")" + mv "$(basename "$ucxdestdir")" "$ucxdestdir" + if test "${ucxversion}" \< "1.17.1"; then + cmd='s/\(#include \)/\1\n#include /' + sed -i.orig "$cmd" "$ucxdestdir/src/ucs/time/time.h" + fi + else + echo reusing directory "$ucxdestdir"... + fi fi -echo copying license file... + +if test "$mpiname" = "mpich"; then + mpidate=$(sed -nE "s/MPICH_RELEASE_DATE=\"(.*)\"/\1/p" "$SOURCE/configure") +fi +if test "$mpiname" = "openmpi"; then + mpidate=$(sed -nE "s/date=\"(.*)\"/\1/p" "$SOURCE/VERSION") +fi +if test -n "${mpidate+x}"; then + case "$(uname)" in + Linux) + timestamp=$(date -d "$mpidate" "+%s") + ;; + Darwin) + datefmt="%b %d, %Y %T%z" + if test "$mpiname" = "mpich"; then + mpidate=$(awk '{$3=$3",";print $2,$3,$NF}' <<< "$mpidate") + fi + timestamp=$(date -j -f "$datefmt" "$mpidate 12:00:00+0000" "+%s") + ;; + esac + echo writing source-date-epoch ... + echo "$timestamp" > "$SOURCE/source-date-epoch" +fi + +echo writing package metadata ... +echo "Name: $mpiname" > "$PACKAGE/METADATA" +echo "Version: $version" >> "$PACKAGE/METADATA" + +echo copying MPI license file... cp "$SOURCE"/"$license" "$PACKAGE/LICENSE" +if test -n "${ofidestdir+x}"; then + echo copying OFI license file... + cp "$ofidestdir/COPYING" "$PACKAGE/LICENSE.ofi" +fi +if test -n "${ucxdestdir+x}"; then + echo copying UCX license file... + cp "$ucxdestdir/LICENSE" "$PACKAGE/LICENSE.ucx" +fi +if test "$mpiname" = "mpich"; then + licenses=( + "$SOURCE"/modules/hwloc/COPYING + "$SOURCE"/modules/json-c/COPYING + "$SOURCE"/modules/yaksa/COPYRIGHT + ) +fi +if test "$mpiname" = "openmpi"; then + licenses=( + "$SOURCE"/3rd-party/hwloc-*/COPYING + "$SOURCE"/3rd-party/libevent-*/LICENSE + "$SOURCE"/3rd-party/openpmix/LICENSE + "$SOURCE"/3rd-party/prrte/LICENSE + "$SOURCE"/3rd-party/treematch/COPYING + ) +fi +for license in "${licenses[@]}"; do + test -f "$license" || continue + module=$(basename "$(dirname "$license")") + module="${module%%-[0-9]*}" + echo copying "$module" license file... + cp "$license" "$PACKAGE/LICENSE.$module" +done diff --git a/build-wheel.sh b/build-wheel.sh index 89d5c44..002308a 100755 --- a/build-wheel.sh +++ b/build-wheel.sh @@ -6,6 +6,9 @@ WORKDIR=package/workdir DESTDIR=package/install ARCHLIST=${ARCHLIST:-$(uname -m)} +read -r SOURCE_DATE_EPOCH < "$SOURCE"/source-date-epoch +export SOURCE_DATE_EPOCH=$SOURCE_DATE_EPOCH + export CIBW_BUILD_FRONTEND='build' export CIBW_BUILD='cp313-*' export CIBW_SKIP='*musllinux*' @@ -19,8 +22,8 @@ if test "$(uname)" = Linux; then containerengine=$(basename "$(command -v podman || command -v docker)") export CIBW_CONTAINER_ENGINE=$containerengine export SOURCE="/project/$SOURCE" - export WORKDIR="/host/$PWD/$WORKDIR" - export DESTDIR="/host/$PWD/$DESTDIR" + export WORKDIR="/host$PWD/$WORKDIR" + export DESTDIR="/host$PWD/$DESTDIR" platform=linux fi if test "$(uname)" = Darwin; then @@ -31,7 +34,7 @@ if test "$(uname)" = Darwin; then platform=macos fi -python -m pipx run \ +pipx run \ cibuildwheel \ --platform "$platform" \ --output-dir "${1:-dist}" \ diff --git a/check-wheel.sh b/check-wheel.sh index 0b7e582..8bdb6f4 100755 --- a/check-wheel.sh +++ b/check-wheel.sh @@ -16,7 +16,9 @@ unzip -qq "$wheelfile" -d "$workdir" cd "$workdir" whlname=$(basename "$wheelfile") -pkgname=${whlname%%-*} +whlinfo=${whlname%%-py*} +pkgname=${whlinfo%%-*} +version=${whlinfo##*-} mpiname=${pkgname} data=$(ls -d "$pkgname"-*.data/data) @@ -29,9 +31,12 @@ if test "$(uname)" = Linux; then libsdir=.libs print-runpath() { patchelf --print-rpath "$1"; } print-needed() { patchelf --print-needed "$1"; } - if test -f "$data"/lib/libucp.so; then + if test -f "$data"/lib/*/libucp.so.*; then runlibs=$runlibs'|libuc(m|p|s|t)'$soregex fi + if test -f "$data"/lib/*/libfabric.so.*; then + runlibs=$runlibs'|libfabric'$soregex + fi fi if test "$(uname)" = Darwin; then runpath='@executable_path/../lib/|@loader_path/' @@ -41,9 +46,12 @@ if test "$(uname)" = Darwin; then libsdir=.dylibs print-runpath() { otool -l "$1" | sed -n '/RPATH/{n;n;p;}'; } print-needed() { otool -L "$1" | sed 1,1d; } - if test -f "$data"/lib/libucp.dylib; then + if test -f "$data"/lib/libucp.*.dylib; then runlibs=$runlibs'|libuc(m|p|s|t)'$soregex fi + if test -f "$data"/lib/libfabric.*.dylib; then + runlibs=$runlibs'|libfabric'$soregex + fi fi if test "$mpiname" = "mpich"; then @@ -65,12 +73,17 @@ if test "$mpiname" = "mpich"; then "$data"/bin/hydra_* ) libraries=( - "$data"/lib/libmpi.* + "$data"/lib/lib*mpi.* ) - if test -d "$data"/lib/ucx; then + if ls "$data"/lib/*/libfabric.* > /dev/null 2>&1; then + libraries+=( + "$data"/lib/*/libfabric.* + ) + fi + if ls "$data"/lib/*/libucp.* > /dev/null 2>&1; then libraries+=( - "$data"/lib/libuc[mpst]*.* - "$data"/lib/ucx/libuc*.* + "$data"/lib/*/libuc[mpst]*.* + "$data"/lib/*/ucx/libuc[mpst]*.* ) fi fi @@ -96,6 +109,25 @@ if test "$mpiname" = "openmpi"; then "$data"/lib/libmpi.* "$data"/lib/libopen-*.* ) + if test "${version%%.*}" -ge 5; then + libraries+=( + "$data"/lib/libevent*.* + "$data"/lib/libhwloc.* + "$data"/lib/libpmix.* + "$data"/lib/libprrte.* + ) + fi + if ls "$data"/lib/*/libucp.* > /dev/null 2>&1; then + libraries+=( + "$data"/lib/*/libuc[mpst]*.* + "$data"/lib/*/ucx/libuc[mpst]*.* + ) + fi + if ls "$data"/lib/*/libfabric.* > /dev/null 2>&1; then + libraries+=( + "$data"/lib/*/libfabric.* + ) + fi runlibs+='|lib(z|util|event.*|hwloc)'$soregex runlibs+='|lib(open-(pal|rte)|pmix|prrte)'$soregex fi @@ -103,18 +135,19 @@ fi check-binary() { local dso=$1 out1="" out2="" echo checking "$dso"... - test -f "$dso" || printf "ERROR: file not found\n" + test -f "$dso" || (printf "ERROR: file not found\n"; exit 1) out1="$(print-runpath "$dso" | grep -vE "$runpath" || true)" test -z "$out1" || printf "ERROR: RUNPATH\n%s\n" "$out1" out2="$(print-needed "$dso" | grep -vE "$runlibs" || true)" test -z "$out2" || printf "ERROR: NEEDED\n%s\n" "$out2" - test -z "$out1" -a -z "$out2" + test -z "$out1" + test -z "$out2" } for header in "${headers[@]-}"; do test -n "$header" || break echo checking "$header"... - test -f "$header" + test -f "$header" || (printf "ERROR: file not found\n"; exit 1) out=$(grep -E '^#include\s+"mpicxx\.h"' "$header" || true) test -z "$out" || printf "ERROR: include\n%s\n" "$out" test -z "$out" @@ -122,7 +155,7 @@ done for script in "${scripts[@]-}"; do test -n "$script" || break echo checking "$script"... - test -f "$script" + test -f "$script" || (printf "ERROR: file not found\n"; exit 1) out=$(grep -E "/opt/$mpiname" "$script" || true) test -z "$out" || printf "ERROR: prefix\n%s\n" "$out" test -z "$out" diff --git a/cibw-build-mpi.sh b/cibw-build-mpi.sh index f68b325..16a27b0 100755 --- a/cibw-build-mpi.sh +++ b/cibw-build-mpi.sh @@ -1,47 +1,138 @@ #!/bin/bash set -euo pipefail -mpiname="${MPINAME:-mpich}" - PROJECT=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) PACKAGE=$PROJECT/package SOURCE=${SOURCE:-$PACKAGE/source} WORKDIR=${WORKDIR:-$PACKAGE/workdir} DESTDIR=${DESTDIR:-$PACKAGE/install} + +test -f "$PACKAGE"/METADATA +mpiname=$(awk '/Name/{print $2}' "$PACKAGE"/METADATA) +version=$(awk '/Version/{print $2}' "$PACKAGE"/METADATA) + +case $(uname) in + Linux) njobs=$(nproc);; + Darwin) njobs=$(sysctl -n hw.physicalcpu);; +esac + +if test "$(uname)" = Darwin; then + export MACOSX_DEPLOYMENT_TARGET="11.0" + if test "$(uname -m)" = x86_64; then + export MACOSX_DEPLOYMENT_TARGET="10.9" + export ac_cv_func_aligned_alloc="no" # macOS>=10.15 + fi + if test -n "${SOURCE_DATE_EPOCH+x}"; then + export ZERO_AR_DATE="YES" + fi +fi + PREFIX=${PREFIX:-"/opt/$mpiname"} +build_cflags+=("-ffile-prefix-map=$SOURCE=/$mpiname-$version") +build_cflags+=("-ffile-prefix-map=$WORKDIR=/$mpiname-$version") +if test "$(uname)" = Linux; then + build_cflags+=("-Wl,--build-id=none") +fi + +case "$mpiname" in + mpich) MODSOURCE="$SOURCE"/modules ;; + openmpi) MODSOURCE="$SOURCE"/3rd-party ;; +esac + +if test -d "$MODSOURCE"/ucx-*; then + ucxdir=$(basename "$(ls -d "$MODSOURCE"/ucx-*)") + UCXSOURCE="$MODSOURCE"/"$ucxdir" + UCXWORKDIR="$WORKDIR"/"$(basename "$MODSOURCE")"/"$ucxdir" + mkdir -p "$UCXWORKDIR" && cd "$UCXWORKDIR" + if test ! -e "$UCXWORKDIR"/config.log; then + echo running configure for UCX + ucxconfigure="$UCXSOURCE"/contrib/configure-release-mt + ucxoptions=(--prefix="$PREFIX" --disable-static) + "$ucxconfigure" "${ucxoptions[@]}" CFLAGS="${build_cflags[*]}" + fi + echo running make with "${njobs:-1}" jobs for UCX + make -j "${njobs:-1}" install-strip DESTDIR="$DESTDIR" + for lib in "${DESTDIR}${PREFIX}"/lib/libuc[mpst].la; do + sed -i "/dependency_libs\s*=/s|\($PREFIX\)|$DESTDIR\1|g" "$lib" + done + for pkg in "${DESTDIR}${PREFIX}"/lib/pkgconfig/ucx*.pc; do + sed -i "/prefix\s*=/s|\($PREFIX\)|$DESTDIR\1|g" "$pkg" + done +fi + +if test -d "$MODSOURCE"/libfabric-*; then + ofidir=$(basename "$(ls -d "$MODSOURCE"/libfabric-*)") + OFISOURCE="$MODSOURCE"/"$ofidir" + OFIWORKDIR="$WORKDIR"/"$(basename "$MODSOURCE")"/"$ofidir" + mkdir -p "$OFIWORKDIR" && cd "$OFIWORKDIR" + if test ! -e "$OFIWORKDIR"/config.log; then + echo running configure for OFI + oficonfigure="$OFISOURCE"/configure + ofioptions=(--prefix="$PREFIX" --disable-static) + "$oficonfigure" "${ofioptions[@]}" CFLAGS="${build_cflags[*]}" + fi + echo running make with "${njobs:-1}" jobs for OFI + make -j "${njobs:-1}" install-strip DESTDIR="$DESTDIR" + for lib in "${DESTDIR}${PREFIX}"/lib/libfabric*.la; do + sed -i "/dependency_libs\s*=/s|\($PREFIX\)|$DESTDIR\1|g" "$lib" + done + for pkg in "${DESTDIR}${PREFIX}"/lib/pkgconfig/libfabric*.pc; do + sed -i "/prefix\s*=/s|\($PREFIX\)|$DESTDIR\1|g" "$pkg" + done +fi + if test "$mpiname" = "mpich"; then - version=$(sed -n 's/MPICH_VERSION=\(.*\)/\1/p' "$SOURCE"/maint/Version) + case $(uname)-$(uname -m) in + Linux-x86_64) netmod=ofi,ucx ;; + Linux-aarch64) netmod=ofi,ucx ;; + Linux-ppc64le) netmod=ofi,ucx ;; + *) netmod=ofi ;; + esac options=( CC=cc CXX=c++ --prefix="$PREFIX" - --with-device=ch4:"${variant:-ofi}" + --with-device="ch4:$netmod" --with-pm=hydra:gforker - --with-libfabric=embedded --with-ucx=embedded + --with-libfabric=embedded --with-hwloc=embedded --with-yaksa=embedded --disable-cxx --disable-doc + --disable-debug + --disable-dlopen --disable-static --disable-opencl --disable-libxml2 --disable-dependency-tracking ) + if test -d "$SOURCE"/modules/ucx-*; then + options=("${options[@]/--with-ucx=embedded/--with-ucx=$DESTDIR$PREFIX}") + fi + if test -d "$SOURCE"/modules/libfabric-*; then + options=("${options[@]/--with-libfabric=embedded/--with-libfabric=$DESTDIR$PREFIX}") + fi + generated_files+=(src/include/mpichinfo.h) + generated_files+=(src/pm/hydra/hydra_config.h) + generated_files+=(src/pm/hydra/include/hydra_config.h) + generated_files+=(modules/ucx/config.h) + generated_files+=(modules/ucx/src/tools/info/build_config.h) + export MPICHLIB_CFLAGS="${build_cflags[*]}" if test "${version%%.*}" -lt 4; then options=("${options[@]/--disable-cxx}") - export FCFLAGS=-fallow-argument-mismatch + options=("${options[@]}" --disable-numa) + export CFLAGS=$MPICHLIB_CFLAGS export FFLAGS=-fallow-argument-mismatch + export FCFLAGS=-fallow-argument-mismatch fi if test "$(uname)" = Darwin; then export MPICH_MPICC_LDFLAGS="-Wl,-rpath,$PREFIX/lib" export MPICH_MPICXX_LDFLAGS="-Wl,-rpath,$PREFIX/lib" export MPICH_MPIFORT_LDFLAGS="-Wl,-rpath,$PREFIX/lib" fi - disable_doc='s/^\(install-data-local:\s\+\)\$/\1#\$/' - sed -i.orig "$disable_doc" "$SOURCE"/Makefile.in fi if test "$mpiname" = "openmpi"; then @@ -49,55 +140,148 @@ if test "$mpiname" = "openmpi"; then CC=cc CXX=c++ --prefix="$PREFIX" - --disable-dlopen - --disable-oshmem - --without-ofi --without-ucx - --without-psm2 + --without-ofi --without-cuda --without-rocm + --with-hwloc=internal + --with-libevent=internal --with-pmix=internal --with-prrte=internal - --with-libevent=internal - --with-hwloc=internal + --enable-ipv6 + --disable-doc + --disable-debug + --disable-dlopen --disable-static + --disable-oshmem --disable-opencl --disable-libxml2 --disable-libompitrace --enable-mpi-fortran=mpifh --disable-dependency-tracking ) + if test -d "$SOURCE"/3rd-party/ucx-*; then + options=("${options[@]/--without-ucx/--with-ucx=$DESTDIR$PREFIX}") + fi + if test -d "$SOURCE"/3rd-party/libfabric-*; then + options=("${options[@]/--without-ofi/--with-ofi=$DESTDIR$PREFIX}") + fi + generated_files+=(ompi/tools/ompi_info/Makefile) + generated_files+=(oshmem/tools/oshmem_info/Makefile) + generated_files+=(3rd-party/openpmix/src/tools/*/Makefile) + generated_files+=(3rd-party/prrte/src/tools/*/Makefile) + generated_files+=(opal/include/opal_config.h) + generated_files+=(opal/include/opal/version.h) + generated_files+=(3rd-party/openpmix/src/include/pmix_config.h) + generated_files+=(3rd-party/prrte/src/include/prte_config.h) + generated_files+=(ompi/tools/mpisync/Makefile) + generated_files+=(orte/tools/orte-info/Makefile) + generated_files+=(opal/mca/pmix/pmix3x/pmix/src/tools/pmix_info/Makefile) + generated_files+=(opal/mca/pmix/pmix3x/pmix/src/include/pmix_config.h) + export CFLAGS="${build_cflags[*]}" + export USER=user HOSTNAME=localhost fi -if test "$(uname)" = Darwin; then - export MACOSX_DEPLOYMENT_TARGET="11.0" - if test "$(uname -m)" = x86_64; then - export MACOSX_DEPLOYMENT_TARGET="10.9" - export ac_cv_func_aligned_alloc="no" # macOS>=10.15 - fi +mkdir -p "$WORKDIR" && cd "$WORKDIR" + +if test ! -e "$WORKDIR"/config.log; then + echo running configure for MPI + "$SOURCE"/configure "${options[@]}" || cat config.log + # shellcheck disable=SC2206 + generated_files=(${generated_files[@]:-}) + for filename in "${generated_files[@]}"; do + test -n "$filename" || continue + test -f "$filename" || continue + source="s|$SOURCE|/$mpiname-$version|g" + workdir="s|$WORKDIR|/$mpiname-$version|g" + destdir="s|$DESTDIR||g" + echo removing source/build/install paths in "$filename" + if test "$(basename "$filename")" = "Makefile"; then + sed -i.orig "/-D.*_BUILD_CFLAGS=/$source" "$filename" + sed -i.orig "/-D.*_BUILD_CFLAGS=/$workdir" "$filename" + sed -i.orig "/-D.*_BUILD_CFLAGS=/$destdir" "$filename" + sed -i.orig "/-D.*_BUILD_CPPFLAGS=/$source" "$filename" + sed -i.orig "/-D.*_BUILD_CPPFLAGS=/$workdir" "$filename" + sed -i.orig "/-D.*_BUILD_CPPFLAGS=/$destdir" "$filename" + sed -i.orig "/-D.*_BUILD_LIBS=/$source" "$filename" + sed -i.orig "/-D.*_BUILD_LIBS=/$workdir" "$filename" + sed -i.orig "/-D.*_BUILD_LIBS=/$destdir" "$filename" + else + sed -i.orig "$source" "$filename" + sed -i.orig "$workdir" "$filename" + sed -i.orig "$destdir" "$filename" + fi + done fi -case $(uname) in - Linux) njobs=$(nproc);; - Darwin) njobs=$(sysctl -n hw.physicalcpu);; -esac +echo running make with "${njobs:-1}" jobs for MPI +make -j "${njobs:-1}" install-strip DESTDIR="$DESTDIR" + +fixup-ucx() { + +cd "${DESTDIR}${PREFIX}" +rm -fr include/uc[mpst] +rm -f bin/io_demo +rm -f bin/ucx_* +rm -f lib/libuc[mpst]*.a +rm -f lib/libuc[mpst]*.la +rm -f lib/ucx/libuc[mt]_*.a +rm -f lib/ucx/libuc[mt]_*.la +rm -fr lib/cmake/ucx +rm -f lib/pkgconfig/ucx*.pc +rm -fr share/ucx + +cd "${DESTDIR}${PREFIX}/lib" +test -f libucp.so || return 0 +for lib in libuc[mpst]*.so.?; do + if test -f "$lib".*; then + mv "$(readlink "$lib")" "$lib" + ln -sf "$lib" "${lib%.*}" + fi + patchelf --set-rpath "\$ORIGIN" "$lib" +done +patchelf --add-rpath "\$ORIGIN/ucx" libucm.so.? +patchelf --add-rpath "\$ORIGIN/ucx" libuct.so.? +for lib in ucx/libuc[mt]_*.so.?; do + if test -f "$lib".*; then + mv "$(dirname "$lib")/$(readlink "$lib")" "$lib" + ln -srf "$lib" "${lib%.*}" + fi + patchelf --set-rpath "\$ORIGIN" "$lib" + patchelf --add-rpath "\$ORIGIN/.." "$lib" +done + +} # fixup-ucx() -mkdir -p "$WORKDIR" -cd "$WORKDIR" +fixup-ofi() { -echo running configure -"$SOURCE"/configure "${options[@]}" || cat config.log +cd "${DESTDIR}${PREFIX}" +rm -fr include/rdma +rm -f bin/fi_* +rm -f lib/libfabric.a +rm -f lib/libfabric.la +rm -f lib/pkgconfig/libfabric.pc +rm -f share/man/man?/fabric.? +rm -f share/man/man?/fi_*.? + +cd "${DESTDIR}${PREFIX}/lib" +test -f libfabric.so || return 0 +for lib in libfabric.so.?; do + if test -f "$lib".*.*; then + mv "$(readlink "$lib")" "$lib" + ln -sf "$lib" "${lib%.*}" + fi + patchelf --set-rpath "\$ORIGIN" "$lib" +done -echo running make with "${njobs:-1}" jobs -make -j "${njobs:-1}" install DESTDIR="$DESTDIR" +} # fixup-ofi() -fixup-mpich() { +fixup-mpi-mpich() { cd "${DESTDIR}${PREFIX}" rm -f include/*cxx.h rm -f include/*.mod rm -f include/*f.h -rm -fr include/rdma rm -f bin/mpif77 rm -f bin/mpif90 rm -f bin/mpifort @@ -112,13 +296,6 @@ rm -f lib/lib*mpifort.* rm -fr lib/pkgconfig rm -fr share -cd "${DESTDIR}${PREFIX}" -rm -f bin/io_demo -rm -f bin/ucx_read_profile -rm -f lib/libuc[mpst]*.la -rm -f lib/ucx/libuct_*.la -rm -fr lib/cmake - headers=(mpi.h) scripts=(mpicc mpic++ mpicxx) executables=(mpichversion mpivars) @@ -135,7 +312,11 @@ for script in "${scripts[@]}"; do # shellcheck disable=SC2016 topdir='$(CDPATH= cd -- "$(dirname -- "${BASH_SOURCE:-$0}")/.." \&\& pwd)' sed -i.orig s@^prefix=.*@prefix="$topdir"@ "$script" + sed -i.orig s:"\s*${build_cflags[*]}"::g "$script" sed -i.orig s:"$PREFIX":\"\$\{prefix\}\":g "$script" + sed -i.orig s:"$WORKDIR.*/lib.*\.la"::g "$script" + sed -i.orig s:"$DESTDIR.*/lib.*\.la"::g "$script" + sed -i.orig s:"$DESTDIR"::g "$script" sed -i.orig s:-Wl,-commons,use_dylibs::g "$script" sed -i.orig s:/usr/bin/bash:/bin/bash:g "$script" sed -i.orig s:-lmpicxx::g "$script" @@ -143,37 +324,42 @@ for script in "${scripts[@]}"; do done if test "$(uname)" = Linux; then - libmpi="libmpi.so.12" cd "${DESTDIR}${PREFIX}/bin" for exe in "${executables[@]}"; do patchelf --set-rpath "\$ORIGIN/../lib" "$exe" done cd "${DESTDIR}${PREFIX}/lib" - if test -f "$libmpi".*.*; then - mv "$(readlink "$libmpi")" "$libmpi" - ln -sf "$libmpi" "${libmpi%.*}" + for lib in libfabric.so.? libuc[mpst]*.so.?; do + test -f "$lib" || continue + for exe in "${executables[@]}"; do + patchelf --remove-needed "$lib" "../bin/$exe" + done + done + cd "${DESTDIR}${PREFIX}/lib" + rm -fr "$mpiname" + if test -f libfabric.so; then + mkdir -p "$mpiname" + mv libfabric.* "$mpiname" fi if test -f libucp.so; then - patchelf --set-rpath "\$ORIGIN" "$libmpi" - for lib in libuc[mpst]*.so.?; do - if test -f "$lib".*; then - mv "$(readlink "$lib")" "$lib" - ln -sf "$lib" "${lib%.*}" - fi - patchelf --set-rpath "\$ORIGIN" "$lib" - for exe in "${executables[@]}"; do - patchelf --remove-needed "$lib" "../bin/$exe" - done - done - patchelf --add-rpath "\$ORIGIN/ucx" libuct.so.? - for lib in ucx/libuct_*.so.?; do - if test -f "$lib".*; then - mv "$(dirname "$lib")/$(readlink "$lib")" "$lib" - ln -srf "$lib" "${lib%.*}" - fi - patchelf --set-rpath "\$ORIGIN/.." "$lib" - done + mkdir -p "$mpiname" + mv libuc[mpst]*.* ucx "$mpiname" fi + cd "${DESTDIR}${PREFIX}/lib" + for lib in lib*.so; do + patchelf --set-rpath "\$ORIGIN" "$lib" + soname=$(patchelf --print-soname "$lib") + if test -L "$soname"; then + mv "$(readlink "$soname")" "$soname" + ln -sf "$soname" "$lib" + fi + done + if test -d "$mpiname"; then + patchelf --add-rpath "\$ORIGIN/$mpiname" libmpi.so.* + fi + cd "${DESTDIR}${PREFIX}/lib" + find . -name '*.so' -type l -delete + ln -s libmpi.so.* libmpi.so fi if test "$(uname)" = Darwin; then @@ -182,9 +368,10 @@ if test "$(uname)" = Darwin; then libpmpi="libpmpi.12.dylib" cd "${DESTDIR}${PREFIX}/bin" for exe in "${executables[@]}"; do - install_name_tool -change "$libdir/$libmpi" "@rpath/$libmpi" "$exe" - install_name_tool -change "$libdir/$libpmpi" "@rpath/$libpmpi" "$exe" install_name_tool -add_rpath "@executable_path/../lib/" "$exe" + for lib in "$libmpi" "$libpmpi"; do + install_name_tool -change "$libdir/$lib" "@rpath/$lib" "$exe" + done done cd "${DESTDIR}${PREFIX}/lib" for lib in "$libmpi" "$libpmpi"; do @@ -192,11 +379,15 @@ if test "$(uname)" = Darwin; then install_name_tool -add_rpath "@loader_path/" "$lib" done install_name_tool -change "$libdir/$libpmpi" "@rpath/$libpmpi" "$libmpi" + cd "${DESTDIR}${PREFIX}/lib" + find . -name '*.dylib' -type l -delete + ln -s "$libmpi" libmpi.dylib + ln -s "$libpmpi" libpmpi.dylib fi -} # fixup-mpich() +} # fixup-mpi-mpich() -fixup-openmpi() { +fixup-mpi-openmpi() { cd "${DESTDIR}${PREFIX}" rm -fr include/ev* @@ -227,6 +418,7 @@ rm -f bin/psched rm -f bin/pterm rm -fr sbin +rm -f bin/pcc rm -f bin/pmixcc rm -f bin/mpif77 rm -f bin/mpif90 @@ -235,6 +427,8 @@ rm -f bin/oshrun rm -f lib/*.mod rm -f lib/lib*.a rm -f lib/lib*.la +rm -f lib/*/lib*.a +rm -f lib/*/lib*.la rm -f lib/libmpi_mpifh*.* rm -f lib/libmpi_usempi*.* rm -f lib/libompitrace.* @@ -245,6 +439,7 @@ rm -fr share/doc rm -fr share/man rm -fr share/hwloc rm -fr share/prte/rst +rm -fr share/ucx rm -f share/openmpi/mpif77-wrapper-data.txt rm -f share/openmpi/mpif90-wrapper-data.txt rm -f share/openmpi/mpifort-wrapper-data.txt @@ -274,6 +469,23 @@ if test "$(uname)" = Linux; then patchelf --set-rpath "\$ORIGIN/../lib" "$exe" done cd "${DESTDIR}${PREFIX}/lib" + for lib in libfabric.so.? libuc[mpst]*.so.?; do + test -f "$lib" || continue + for exe in "${executables[@]}"; do + patchelf --remove-needed "$lib" "../bin/$exe" + done + done + cd "${DESTDIR}${PREFIX}/lib" + rm -fr "$mpiname" + if test -f libfabric.so; then + mkdir -p "$mpiname" + mv libfabric.* "$mpiname" + fi + if test -f libucp.so; then + mkdir -p "$mpiname" + mv libuc[mpst]*.* ucx "$mpiname" + fi + cd "${DESTDIR}${PREFIX}/lib" for lib in lib*.so; do patchelf --set-rpath "\$ORIGIN" "$lib" soname=$(patchelf --print-soname "$lib") @@ -282,6 +494,13 @@ if test "$(uname)" = Linux; then ln -sf "$soname" "$lib" fi done + if test -d "$mpiname"; then + patchelf --add-rpath "\$ORIGIN/$mpiname" libmpi.so.* + patchelf --add-rpath "\$ORIGIN/$mpiname" libopen-*.so.* + fi + cd "${DESTDIR}${PREFIX}/lib" + find . -name '*.so' -type l -delete + ln -s libmpi.so.* libmpi.so fi if test "$(uname)" = Darwin; then @@ -303,6 +522,7 @@ if test "$(uname)" = Darwin; then for lib in lib*.*.dylib; do install_name_tool -id "@rpath/$lib" "$lib" install_name_tool -add_rpath "@loader_path/" "$lib" + unset dependencies while IFS= read -r dep do dependencies+=("$dep") done < <(otool -L "$lib" | awk '(NR>1) {print $1}') @@ -313,6 +533,9 @@ if test "$(uname)" = Darwin; then fi done done + cd "${DESTDIR}${PREFIX}/lib" + find . -name '*.dylib' -type l -delete + ln -s libmpi.*.dylib libmpi.dylib fi cd "${DESTDIR}${PREFIX}/share/openmpi" @@ -326,6 +549,8 @@ for cmd in mpicc mpic++ mpicxx mpiCC ortecc; do flags="\(linker_flags=-L\${libdir}\) *" rpath="-Wl,-rpath,\${libdir}" sed -i.orig "s:^$flags$:\1 $rpath:" "$wrapper_data" + libdir="$DESTDIR$PREFIX/lib" + sed -i.orig "s:$libdir:\${libdir}:g" "$wrapper_data" rm "$wrapper_data".orig done @@ -333,7 +558,7 @@ cd "${DESTDIR}${PREFIX}/bin" wrapper_cmd=opal_wrapper wrapper_src="$PROJECT/cibw-ompi-wrapper.c" wrapper_bin="$WORKDIR/cibw-ompi-wrapper.exe" -cc -DWRAPPER="$wrapper_cmd" "$wrapper_src" -o "$wrapper_bin" +cc -O -DWRAPPER="$wrapper_cmd" "$wrapper_src" -o "$wrapper_bin" executables=(mpicc mpic++ mpicxx mpiCC ortecc) for exe in "${executables[@]}"; do test -e "$exe" || continue @@ -353,15 +578,30 @@ mv "$executor" "$wrapper" ln -s "$wrapper" "$executor" wrapper_src="$PROJECT/cibw-ompi-wrapper.c" wrapper_bin="$WORKDIR/cibw-ompi-wrapper.exe" -cc -DWRAPPER="$wrapper" "$wrapper_src" -o "$wrapper_bin" +cc -O -DWRAPPER="$wrapper" "$wrapper_src" -o "$wrapper_bin" executables=(mpirun mpiexec orterun) for exe in "${executables[@]}"; do test -e "$exe" || continue test -L "$exe" || continue - install "$wrapper_bin" "$exe" + install -s "$wrapper_bin" "$exe" done -} # fixup-openmpi() +} # fixup-mpi-openmpi() + +echo fixing UCX install tree +fixup-ucx +echo fixing OFI install tree +fixup-ofi +echo fixing MPI install tree +fixup-mpi-"$mpiname" -echo fixing install tree -fixup-"$mpiname" +echo checking install tree +cd "${DESTDIR}${PREFIX}" +dirty=0 +echo checking for files with SOURCE +if grep -lr "$SOURCE" ; then dirty=1; fi +echo checking for files with WORKDIR +if grep -lr "$WORKDIR"; then dirty=1; fi +echo checking for files with DESTDIR +if grep -lr "$DESTDIR"; then dirty=1; fi +test "$dirty" -eq 0 diff --git a/cibw-check-mpi.sh b/cibw-check-mpi.sh index 0a9f87b..b27326e 100755 --- a/cibw-check-mpi.sh +++ b/cibw-check-mpi.sh @@ -20,7 +20,11 @@ int main(int argc, char *argv[]) MPI_Comm_size(MPI_COMM_WORLD, &size); MPI_Comm_rank(MPI_COMM_WORLD, &rank); MPI_Get_processor_name(name, &len); + if (rank != 0) + MPI_Recv(name, 0 , MPI_BYTE, rank -1, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE); printf("Hello, World! I am process %d of %d on %s.\n", rank, size, name); + if (rank != size - 1) + MPI_Send(name, 0 , MPI_BYTE, rank +1, 0, MPI_COMM_WORLD); MPI_Finalize(); return 0; } @@ -68,3 +72,25 @@ RUN mpicxx helloworld.cxx -o helloworld-cxx RUN command -v mpiexec RUN mpiexec -n 3 ./helloworld-c RUN mpiexec -n 3 ./helloworld-cxx + +if test "$mpiname" = "mpich"; then + version=$(mpichversion --version | cut -d':' -f 2) + case $(uname)-$(uname -m) in + Linux-x86_64) ch4netmods=(ofi ucx) ;; + Linux-aarch64) ch4netmods=(ofi ucx) ;; + Linux-ppc64le) ch4netmods=(ofi) ;; # FIXME + *) ch4netmods=(ofi) ;; + esac + export MPICH_CH4_OFI_CAPABILITY_DEBUG=1 + export MPICH_CH4_UCX_CAPABILITY_DEBUG=1 + for netmod in "${ch4netmods[@]}"; do + printf "testing ch4:%s ... " "$netmod" + export MPICH_CH4_NETMOD="$netmod" + test "${version%%.*}" -ge 5 || netmod="" + ./helloworld-c | grep -i "$netmod" > /dev/null + printf "OK\n" + done + unset MPICH_CH4_OFI_CAPABILITY_DEBUG + unset MPICH_CH4_UCX_CAPABILITY_DEBUG + unset MPICH_CH4_NETMOD +fi diff --git a/cibw-ompi-wrapper.c b/cibw-ompi-wrapper.c index c967ebc..16e72dc 100644 --- a/cibw-ompi-wrapper.c +++ b/cibw-ompi-wrapper.c @@ -34,7 +34,7 @@ int main(int argc, char *argv[]) { (void) readlink("/proc/curproc/file", exe, PATH_MAX); #elif defined(__linux__) (void) readlink("/proc/self/exe", exe, PATH_MAX); -#elif defined( __sun) +#elif defined(__sun) (void) strncpy(exe, getexecname(), PATH_MAX); #else # error unknown system diff --git a/package/pyproject.toml b/package/pyproject.toml index 9787c3b..ac76de3 100644 --- a/package/pyproject.toml +++ b/package/pyproject.toml @@ -1,3 +1,3 @@ [build-system] -requires = ["setuptools", "wheel"] +requires = ["setuptools >= 70.1"] build-backend = "setuptools.build_meta" diff --git a/package/setup.py b/package/setup.py index b515d00..d74b1e0 100644 --- a/package/setup.py +++ b/package/setup.py @@ -1,5 +1,8 @@ from setuptools import setup -from wheel.bdist_wheel import bdist_wheel +try: + from setuptools.command.bdist_wheel import bdist_wheel +except ImportError: + from wheel.bdist_wheel import bdist_wheel import re import os @@ -16,15 +19,18 @@ def get_tag(self): return (self.python_tag, "none", plat_tag) -mpiname = os.environ.get("MPINAME", "mpich") +with open("METADATA") as fobj: + metadata = re.search(r""" + Name:\s*(?P.*)\n + Version:\s*(?P.*)\n + """, fobj.read(), re.VERBOSE).groupdict() + +mpiname = metadata["name"] +version = metadata["version"] release = os.environ.get("RELEASE", "") -pkgname = mpiname if mpiname == "mpich": project = "MPICH" - version_re = re.compile(r""" - \#define\s+MPICH_VERSION\s+\"(.*)\" - """, re.VERBOSE) license = "LicenseRef-MPICH" author = "MPICH Team" author_email = "discuss@mpich.org" @@ -39,11 +45,6 @@ def get_tag(self): if mpiname == "openmpi": project = "Open MPI" - version_re = re.compile(r""" - \#define\s+OMPI_MAJOR_VERSION\s+(\d+)\n - \#define\s+OMPI_MINOR_VERSION\s+(\d+)\n - \#define\s+OMPI_RELEASE_VERSION\s+(\d+)\n - """, re.VERBOSE) license = "LicenseRef-OpenMPI" author = "Open MPI Team" author_email = "users@lists.open-mpi.org" @@ -66,9 +67,6 @@ def get_tag(self): destdir = os.environ.get("DESTDIR", f"{basedir}/install") rootdir = f"{destdir}{prefix}" -with open(f"{rootdir}/include/mpi.h") as fobj: - version = ".".join(version_re.search(fobj.read()).groups()) - data_files = [] for path, dirs, files in os.walk(rootdir): dirs.sort() @@ -82,10 +80,10 @@ def get_tag(self): } setup( - name=pkgname, + name=mpiname, version=version, license=license, - license_files=["LICENSE"], + license_files=["LICENSE", "LICENSE.*"], description=description, long_description=long_description, long_description_content_type="text/x-rst", diff --git a/patches/mpich-4.2.3 b/patches/mpich-4.2.3 new file mode 100644 index 0000000..7abaa41 --- /dev/null +++ b/patches/mpich-4.2.3 @@ -0,0 +1,116 @@ +From 135459663554992589868b9f4e75d040e0cf4942 Mon Sep 17 00:00:00 2001 +From: Ken Raffenetti +Date: Wed, 23 Oct 2024 09:58:33 -0500 +Subject: [PATCH] ch4/ucx: Reorder function definitions + +When building a non-inline netmod configuration, the compiler will flag +MPIDI_UCX_do_am_recv for implicit declaration. Move its definition +before its used to resolve. Fixes pmodels/mpich#7185. +--- + src/mpid/ch4/netmod/ucx/ucx_am.c | 86 ++++++++++++++++---------------- + 1 file changed, 43 insertions(+), 43 deletions(-) + +diff --git a/src/mpid/ch4/netmod/ucx/ucx_am.c b/src/mpid/ch4/netmod/ucx/ucx_am.c +index 082d35f9cc8..05f849eac6d 100644 +--- a/src/mpid/ch4/netmod/ucx/ucx_am.c ++++ b/src/mpid/ch4/netmod/ucx/ucx_am.c +@@ -70,6 +70,49 @@ void MPIDI_UCX_am_send_callback(void *request, ucs_status_t status) + } + + #ifdef HAVE_UCP_AM_NBX ++/* Called when recv buffer is posted */ ++int MPIDI_UCX_do_am_recv(MPIR_Request * rreq) ++{ ++ void *recv_buf; ++ bool is_contig; ++ MPI_Aint data_sz, in_data_sz; ++ int vci = MPIDI_Request_get_vci(rreq); ++ ++ MPIDIG_get_recv_buffer(&recv_buf, &data_sz, &is_contig, &in_data_sz, rreq); ++ if (!is_contig || in_data_sz > data_sz) { ++ /* non-contig datatype, need receive into pack buffer */ ++ /* ucx will error out if buffer size is less than the promised data size, ++ * also use a pack buffer in this case */ ++ recv_buf = MPL_malloc(in_data_sz, MPL_MEM_OTHER); ++ MPIR_Assert(recv_buf); ++ MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = recv_buf; ++ } else { ++ MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = NULL; ++ } ++ ++ MPIDI_UCX_ucp_request_t *ucp_request; ++ size_t received_length; ++ ucp_request_param_t param = { ++ .op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_RECV_INFO, ++ .cb.recv_am = &MPIDI_UCX_am_recv_callback_nbx, ++ .recv_info.length = &received_length, ++ }; ++ void *data_desc = MPIDI_UCX_AM_RECV_REQUEST(rreq, data_desc); ++ /* note: use in_data_sz to match promised data size */ ++ ucp_request = ucp_am_recv_data_nbx(MPIDI_UCX_global.ctx[vci].worker, ++ data_desc, recv_buf, in_data_sz, ¶m); ++ if (ucp_request == NULL) { ++ /* completed immediately */ ++ MPIDI_UCX_ucp_request_t tmp_ucp_request; ++ tmp_ucp_request.req = rreq; ++ MPIDI_UCX_am_recv_callback_nbx(&tmp_ucp_request, UCS_OK, received_length, NULL); ++ } else { ++ ucp_request->req = rreq; ++ } ++ ++ return MPI_SUCCESS; ++} ++ + /* Am handler for messages sent from ucp_am_send_nbx. Registered with + * ucp_worker_set_am_recv_handler. + */ +@@ -116,49 +159,6 @@ ucs_status_t MPIDI_UCX_am_nbx_handler(void *arg, const void *header, size_t head + } + } + +-/* Called when recv buffer is posted */ +-int MPIDI_UCX_do_am_recv(MPIR_Request * rreq) +-{ +- void *recv_buf; +- bool is_contig; +- MPI_Aint data_sz, in_data_sz; +- int vci = MPIDI_Request_get_vci(rreq); +- +- MPIDIG_get_recv_buffer(&recv_buf, &data_sz, &is_contig, &in_data_sz, rreq); +- if (!is_contig || in_data_sz > data_sz) { +- /* non-contig datatype, need receive into pack buffer */ +- /* ucx will error out if buffer size is less than the promised data size, +- * also use a pack buffer in this case */ +- recv_buf = MPL_malloc(in_data_sz, MPL_MEM_OTHER); +- MPIR_Assert(recv_buf); +- MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = recv_buf; +- } else { +- MPIDI_UCX_AM_RECV_REQUEST(rreq, pack_buffer) = NULL; +- } +- +- MPIDI_UCX_ucp_request_t *ucp_request; +- size_t received_length; +- ucp_request_param_t param = { +- .op_attr_mask = UCP_OP_ATTR_FIELD_CALLBACK | UCP_OP_ATTR_FIELD_RECV_INFO, +- .cb.recv_am = &MPIDI_UCX_am_recv_callback_nbx, +- .recv_info.length = &received_length, +- }; +- void *data_desc = MPIDI_UCX_AM_RECV_REQUEST(rreq, data_desc); +- /* note: use in_data_sz to match promised data size */ +- ucp_request = ucp_am_recv_data_nbx(MPIDI_UCX_global.ctx[vci].worker, +- data_desc, recv_buf, in_data_sz, ¶m); +- if (ucp_request == NULL) { +- /* completed immediately */ +- MPIDI_UCX_ucp_request_t tmp_ucp_request; +- tmp_ucp_request.req = rreq; +- MPIDI_UCX_am_recv_callback_nbx(&tmp_ucp_request, UCS_OK, received_length, NULL); +- } else { +- ucp_request->req = rreq; +- } +- +- return MPI_SUCCESS; +-} +- + /* callback for ucp_am_recv_data_nbx */ + void MPIDI_UCX_am_recv_callback_nbx(void *request, ucs_status_t status, size_t length, + void *user_data)