Conda-friendly builds, AMD GPU support, and macOS OpenMP - #23
Merged
Conversation
Squashes the 22 commits of fix-conda-build-and-mpi-detection, rebased onto main so that #15, #16, #17 and #21 are preserved rather than reverted. setup.py - Derive MPI_HOME from the MPI mpi4py is actually linked against, instead of guessing, so the extension and mpi4py cannot disagree at runtime. - Make conda-based builds work: strip the gcc-only tokens RHEL 9 CPython bakes into sysconfig that nvc++ rejects, and drop conda's -B compiler_compat while keeping its -L/-rpath entries. - Make SBD_GPU_ARCH optional and honor it at link as well as compile, so a multi-arch Thrust fatbin keeps every architecture it was asked for. This supersedes the narrower fix in #17. python/__init__.py, device_config.py - Load the backends lazily, one per process, and build all three by default. Removes the silent GPU demotion and the need for a second install. - Report why a backend is missing, and fail instead of silently using the CPU when gpu-omp has no device. README: rewritten around the conda recipe and split GPU prerequisites by build vs run. Retains #15's point that the sdist bundles the headers.
Squashes the four commits of amd-rocm-omp-offload.
The same _core_gpu_omp_offload source now builds under either vendor: nvc++
-mp=gpu for NVIDIA, amdclang++ --offload-arch=gfx* for AMD. Thrust stays
NVIDIA-only, since upstream wires it to nvc++ -cuda and there is no rocThrust
configuration to build; SBD_BUILD_BACKEND=gpu therefore fails on AMD rather
than silently falling back to the CPU.
setup.py
- detect_gpu_toolchain() picks nvidia or amd, with SBD_GPU_VENDOR to force the
choice on a host carrying both; find_rocm_toolchain() honors ROCM_HOME then
PATH, mirroring find_nvidia_hpc_sdk().
- SBD_GPU_ARCH is spelled per vendor and applied at compile and link: nvc++
takes -gpu=cc80,cc90 as one flag, clang needs a repeated --offload-arch=.
Unset on AMD it is detected with the amdgpu-arch that ships beside the chosen
compiler (not one found on PATH, which can come from a different ROCm), and a
GPU-less build host is a hard error rather than an unrunnable module.
- amdclang++ needs far less sysconfig scrubbing than nvc++: only
-fcf-protection, and -march=x86-64-v2 is deliberately left alone.
python/
- 'gpu-omp' is vendor-neutral; get_backend('gpu-omp').__sbd_offload_target__
reports what an install was actually built for, and the Thrust backend is
stamped the same way. rocm/gpu-amd-omp/gpu-rocm-omp are aliases.
- device_config.auto() resolves against the backends actually COMPILED instead
of returning 'gpu' whenever any GPU was detected, which selected the CUDA
Thrust backend on AMD hosts. GPU presence is probed with rocm-smi as well as
nvidia-smi.
- bindings.cpp pins each rank to one offload device, reading the vendor's
device-visibility variable when omp_get_num_devices() reports 0 under dlopen.
README: one rule for MPI -- every GPU backend hands MPI device pointers, so use
a GPU-aware MPI (CUDA-aware on NVIDIA, ROCm-aware on AMD).
Verified on MI250X (gfx90a, ROCm): 8 ranks, 4x2 grid, --device gpu-omp.
sbd_pin_offload_device() was wired into tpb_diag and tpb_diag_from_files but not gdb_diag, so with --device gpu-omp every rank ran gdb_diag on the default device -- the "all ranks land on GPU 0" symptom, for that entry point only. Not caught earlier because the three call sites are copies rather than one helper; folding them together is a separate cleanup.
- Per-probe timeouts in _gpu_available(): rocm-smi keeps the 30 s it needs to enumerate a multi-GCD node, but nvidia-smi no longer inherits it. A wedged NVIDIA driver hangs rather than failing, and it was stalling backend selection for 30 s instead of 5. - Require a GPU[<n>] line from `rocm-smi --showid` rather than trusting its exit status, in both _gpu_available() and DeviceConfig._check_hip(). A ROCm install on a GPU-less BUILD host is exactly where a bare `rocm-smi` exiting 0 claims hardware that is not there, after which auto() picks a backend that cannot run. This is the same signal get_device_info() already counts devices with, so there is now one convention instead of two. - find_rocm_toolchain() honors ROCM_PATH as well as ROCM_HOME, checked in that order so an explicit ROCM_HOME still wins. ROCM_PATH is the variable ROCm itself sets, and tox.ini already passed it through to a build that ignored it.
The Darwin block hardcoded Homebrew paths, so a conda env carrying llvm-openmp still compiled against /opt/homebrew -- while at import time the process loads the conda copies anyway, resolved through the interpreter's rpath. Prefer $CONDA_PREFIX when it has include/omp.h; Homebrew stays as the fallback. Two diagnostics, since this was hard to see: - fail with the fix when neither source has libomp, rather than "'omp.h' file not found" from the middle of a compile - print the resolved compiler and its version. sysconfig records a bare clang++, which PATH resolves, so a Homebrew LLVM silently wins over both Apple clang and a conda toolchain. README: llvm-openmp in both conda recipes, a note to pin CC/CXX, and a troubleshooting entry for a conda compiler package shadowing nvc++/amdclang++. Verified on macOS arm64 under Apple clang 21.0.0 and conda clang 23.1.0: no Homebrew references in the linked extension, pytest and the 2-rank MPI test green in both environments.
Deriving the MPI prefix from mpi4py dropped main's mpicc probe, which broke CI: Debian splits Open MPI's headers to /usr/lib/<triple>/openmpi/include, outside the prefix, so mpi4py names the right MPI but $prefix/include/mpi.h does not exist and the build gave up. A prefix without mpi.h is an unusual layout, not the wrong MPI, so ask mpicc before failing -- both when the prefix lacks the header and when no prefix is found at all. mpi4py stays the preferred source, so the extension and mpi4py still cannot diverge.
--showme:compile is an Open MPI-ism. MPICH's wrapper treats it as a source file and tries to compile it (exit 127 here), so probing that first silently ruled out MPICH -- which this package claims to support. Probe `-show` first, which both wrappers understand and which prints the whole command line, and keep --showme as the second attempt. Duplicate -I/-L entries are collapsed (MPICH repeats them), and the libraries the wrapper names are used as-is: -lmpi for Open MPI, -lmpi -lpmpi for MPICH. Verified against conda-forge openmpi 5.0.10 and mpich 5.0.1.
-show is understood by every wrapper checked -- homebrew open-mpi, conda-forge openmpi 5.0.10, conda-forge mpich 5.0.1 -- while --showme:compile works on the Open MPI ones and exits 127 on MPICH. So -show strictly dominates, and the second attempt was dead weight: one subprocess call instead of up to three.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Supersedes
fix-conda-build-and-mpi-detectionandamd-rocm-omp-offload,rebased onto main so #15, #16, #17 and #21 are preserved.
Build
MPI_HOMEfrom the MPI mpi4py is linked against, so the extension andmpi4py cannot disagree at runtime.
drop conda's
-B compiler_compatwhile keeping its-L/-rpath.SBD_GPU_ARCHis optional and honored at link as well as compile, so amulti-arch Thrust fatbin keeps every architecture (supersedes Pass -gpu= at link time so the Thrust fatbin keeps every architecture #17).
AMD GPUs
amdclang++as well asnvc++;SBD_GPU_VENDORforces the choice on a host with both. Thrust staysNVIDIA-only, and
SBD_BUILD_BACKEND=gpufails on AMD rather than silentlyfalling back to CPU.
amdgpu-archbeside the chosen compiler; aGPU-less build host is a hard error rather than an unrunnable module.
gpu-ompis vendor-neutral;get_backend('gpu-omp').__sbd_offload_target__reports what an install was built for.
device_config.auto()resolves against backends actually compiled, instead ofreturning the CUDA Thrust backend whenever any GPU was detected.
Backends load lazily, one per process, and all are built by default — no
second install, no silent GPU demotion.
macOS takes libomp and OpenBLAS from the conda env (Homebrew as fallback), and
the build prints which compiler and which libomp it chose.
Verified
--device gpu-omp, correct energy.references linked,
pytestand the 2-rank MPI test green.Not exercised: the
gdb_diagoffload-device pin — GDB has no example or testanywhere yet.