diff --git a/README.md b/README.md index ec86a35..71642ba 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,22 @@ # SBD Python Bindings -Python bindings for the Selected Basis Diagonalization (SBD) library with dual CPU/GPU backend support. +Python bindings for the Selected Basis Diagonalization (SBD) library, with CPU and GPU backends. ## Overview -SBD (Selected Basis Diagonalization) is a high-performance library for quantum chemistry calculations. The Python bindings provide access to SBD's **Tensor-Product Basis (TPB)** diagonalization method with support for both CPU and GPU backends. +SBD (Selected Basis Diagonalization) is a high-performance library for quantum chemistry calculations. The Python bindings provide access to SBD's **Tensor-Product Basis (TPB)** diagonalization method on CPU and GPU. **Key Features:** - **TPB diagonalization** for quantum chemistry Hamiltonians -- Dual backend: CPU (OpenMP) and GPU (CUDA), switchable at runtime +- Three backends, selected per call at runtime via `device=`: + `'cpu'` (host OpenMP), `'gpu'` (NVHPC Thrust/CUDA, NVIDIA only) and + `'gpu-omp'` (OpenMP target offload, **NVIDIA or AMD**). All the backends your + toolchain supports can be built into one install; each is imported only when + first used - MPI parallelization - Integration with [qiskit-addon-sqd](https://github.com/Qiskit/qiskit-addon-sqd) for SQD workflows -In addition to TPB, this package also contains experimental support for SBD's **General-Determinant Basis (GDP)** method. However, SBD's **Creation/Annihilation operator (CAOP)** method is currently not supported by this wrapper; users who need it should reference and use the C++ CLI apps in the upstream submodule (`vendor/sbd-upstream/apps/`). +In addition to TPB, this package also contains experimental support for SBD's **General-Determinant Basis (GDB)** method. However, SBD's **Creation/Annihilation operator (CAOP)** method is currently not supported by this wrapper; users who need it should reference and use the C++ CLI apps in the upstream submodule (`vendor/sbd-upstream/apps/`). > [!NOTE] > This package is newly open-sourced. The Python API follows semantic versioning, but the build configuration and GPU backends have been exercised on a limited set of platforms — please report issues. @@ -21,33 +25,56 @@ In addition to TPB, this package also contains experimental support for SBD's ** ### Prerequisites -**Required:** Python 3.10+, MPI (OpenMPI/MPICH), BLAS (OpenBLAS/MKL), pybind11, mpi4py, numpy. +**Required:** Python 3.10+, MPI (OpenMPI/MPICH), BLAS (OpenBLAS/MKL), pybind11, mpi4py, numpy, compiler with OpenMP. -**Optional (GPU):** NVIDIA HPC SDK (nvc++), CUDA-capable GPU, CUDA-aware MPI. +**On macOS:** Apple clang ships without OpenMP, so add `llvm-openmp` to the conda +environment. Homebrew's `libomp` is used as a fallback if the env has none. Pin the +compiler with `CC`/`CXX` as well — a bare `clang++` is resolved through `PATH`, so a +Homebrew LLVM silently wins over both Apple clang and a conda toolchain. The build +prints which compiler and which libomp it chose. + +**For the GPU backends** — optional; without them you get a CPU-only install: + +***On NVIDIA*** (both the Thrust and the OpenMP-offload backend): + +- *to build:* [NVIDIA HPC SDK](https://developer.nvidia.com/hpc-sdk) (`nvc++`). + `SBD_GPU_ARCH` is **optional**: unset, nvc++ targets the GPU of the machine + the toolchain was installed on; set, it is honored exactly and may name + several generations at once (`cc80,cc90,cc100`) — see + [Environment Variables](#environment-variables). + +***On AMD*** (the OpenMP-offload backend only): + +- *to build:* ROCm LLVM toolchain (`amdclang++`). + `SBD_GPU_ARCH` is optional here on a GPU system and is **detected** with + ROCm's `amdgpu-arch` (e.g. `gfx90a` on MI250X, `gfx942` on MI300X). However, on a GPU-less build host, you must + set it. see [Environment Variables](#environment-variables). Either install path compiles the C++ extension on the target machine; -no pre-built wheels are published. The resulting binary depends on -the local MPI and BLAS, so an MPI toolchain (`mpicc` on `PATH`, or -`MPI_HOME` set) and BLAS must be installed first; see -[Environment Variables](#environment-variables). +no pre-built wheels are published. The resulting binary depends on the +local MPI and BLAS, so both must be installed first. ### Install from PyPI -**For a quick CPU test**, a conda environment is the easiest way to satisfy that: -it supplies every dependency — MPI, BLAS and the Python build tools — inside the -environment, so nothing has to be installed system-wide and nothing on the host -is touched. Any conda works; [Miniforge](https://github.com/conda-forge/miniforge) -is a small one that defaults to conda-forge. +**A self-contained conda environment** is the quickest way to get those +dependencies in place for the CPU backend: +```bash +conda create -y -n sbd -c conda-forge \ + python=3.13.12 pybind11 numpy setuptools wheel openblas pyscf pip mpi4py +# ...plus llvm-openmp on macOS +``` ```bash -conda create -n sbd python=3.12 pybind11 numpy setuptools wheel pyscf pip \ - mpi4py openmpi openblas -c conda-forge conda activate sbd +``` + +Now install the sbd-eigensolver-python package +``` pip install sbd-eigensolver ``` The published source distribution (sdist) bundles the sbd header files, so this -needs no git checkout and no submodule step. +needs no git checkout and no submodule step. ### Install from git checkout @@ -69,213 +96,123 @@ etc.), advance the local submodule and rebuild: ```bash git submodule update --remote vendor/sbd-upstream +``` + +After the upstream sbd code is cloned, run: +``` pip install -e . --no-build-isolation --force-reinstall --no-deps ``` ### Environment Variables ```bash -# --- always needed --- -export MPI_HOME=/path/to/mpi -export BLAS_LIB_PATH=/path/to/blas/lib -export BLAS_LIBS=openblas # or mkl_rt - -# --- macOS only — pin the host compiler to system clang so libc++ -# matches Python's. Not needed on Linux GPU boxes (setup.py -# auto-picks nvc++ for the GPU extensions; the CPU extension -# then compiles cleanly under nvc++ as well). -export CC=/usr/bin/clang -export CXX=/usr/bin/clang++ - -# --- GPU backends (Thrust and OpenMP target-offload) — NVIDIA-only. -# Both compile with NVHPC nvc++; there is no AMD/Intel GPU path -# in this wrapper. Requires the NVHPC SDK (one tarball; no -# LLVM/clang source build needed since v1.6). +# --- NVIDIA GPU backends (Thrust and OpenMP target-offload): point at NVHPC. +# Only needed if nvc++ is not already on PATH. Adjust the path. export NVHPC_HOME=/opt/nvidia/hpc_sdk/Linux_x86_64/2025/compilers -# Target GPU arch for nvc++ -gpu=. Used by BOTH the Thrust path -# and the OMP-offload path. nvc++ accepts cc (documented PGI form) -# and sm_. Default cc90 for H100. -# H100: cc90 -# GB200 / B200: cc100 -# A100: cc80 -export SBD_GPU_ARCH=cc100 -``` -You do **not** need to set `CC=nvc` / `CXX=nvc++` or clear `CFLAGS` / -`CXXFLAGS` manually for the GPU builds — setup.py auto-routes those -extensions through nvc++ and filters the RHEL 9 sysconfig flags that -nvc++ rejects. (Earlier versions required this; the v1.6 refactor -moved it into `setup.py:_route_build_through_nvhpc()`.) If you DO -set them, your values win — distutils' `os.environ.setdefault` -semantics. - -The OpenMP target-offload backend cannot share a Python process with -CPU or Thrust. They link incompatible OpenMP runtimes (NVHPC's -`libnvomp` for OMP-offload, `libgomp`/`libomp` for CPU and Thrust), -and `import sbd` eagerly loads every `_core_*.so` it finds in -`python/`. Co-resident `.so` files therefore pull both runtimes into -the same address space, producing "Another OpenMP runtime library has -been detected" and potentially deadlocking at the first `#pragma omp` -region. **The cleanest setup is two venvs with two checkouts** — one -for CPU + Thrust, one for OMP-offload. Within a single venv you can -also switch profiles by removing the other profile's `_core_*.so` -before rebuilding (only files present in `python/` get loaded), but a -second venv avoids the bookkeeping. - -### Build - -Pick **one** of two installation profiles. They produce mutually -incompatible Python processes (different OpenMP runtimes — see the -paragraph above the "Build" section) so put them in **separate venvs -backed by separate checkouts** if you need both. The `source …` -lines below are not optional: forgetting to activate the right venv -before `pip install -e .` either puts the build into the wrong venv -or fails outright. - -> **Note — `setuptools` in the setup line:** because the build uses -> `--no-build-isolation`, it relies on the build tools already present -> in the active venv. Python 3.12+ no longer seeds `setuptools` into -> fresh venvs, so it must be installed explicitly (it is included in -> the `pip install … setuptools wheel` line below). Without it the -> build fails with `Cannot import 'setuptools.build_meta'`. - -**Profile 1 — CPU + Thrust GPU** (the common case) +# --- AMD GPU backend (OpenMP target-offload): point at ROCm LLVM toolchain +# amdclang++ is found on $ROCM_HOME/bin +export ROCM_HOME=/opt/rocm + +# --- OPTIONAL. Only for a host with BOTH GPU toolchains installed, where the +# auto-answer would be an accident of probe order: nvidia | amd | none +export SBD_GPU_VENDOR=amd + +# --- OPTIONAL GPU architecture, spelled per vendor. +# NVIDIA: unset, nvc++ targets the GPU of the machine this toolchain was +# installed on. Set it to pin the target, including SEVERAL at once: +# A100: cc80 H100: cc90 GB200 / B200: cc100 +# ccall-major one target per major generation +# AMD: unset, the arch is DETECTED with ROCm's `amdgpu-arch`. Set it to pin, +# or when building on a host with no GPU (where detection cannot work and +# the build stops asking for it): +# MI250X: gfx90a MI300X: gfx942 (several: gfx90a,gfx942) +export SBD_GPU_ARCH=cc80,cc90,cc100 # NVIDIA +export SBD_GPU_ARCH=gfx90a # AMD + +# --- optional overrides; each has a working default --- +# Which backends to build: defaults to CPU always, plus every GPU backend the +# detected toolchain supports -- Thrust AND OpenMP-offload under nvc++, +# OpenMP-offload only under amdclang++ (there is no rocThrust path). +# Set it only to narrow that: +# cpu CPU only -- skip GPU even if a GPU compiler is present +# gpu Thrust GPU only, no CPU -- NVIDIA only, errors on AMD +# gpu_omp_offload OpenMP target-offload GPU only (either vendor) +export SBD_BUILD_BACKEND=cpu + +# MPI: defaults to whatever mpi4py is linked against. Set this only +# for layouts that cannot be inferred. +# NOTE: Every GPU backend hands MPI device pointers, so use a GPU-aware MPI: +# CUDA-aware on NVIDIA, ROCm-aware on AMD. +export MPI_HOME=/path/to/mpi -```bash -# one-time setup -git clone --recurse-submodules https://github.com/Qiskit/sbd-eigensolver-python.git sbd-thrust -python -m venv ~/venvs/sbd-thrust -source ~/venvs/sbd-thrust/bin/activate -MPICC=$(which mpicc) pip install --no-binary=mpi4py mpi4py pybind11 numpy setuptools wheel - -# build (re-run after pulling main) -source ~/venvs/sbd-thrust/bin/activate # always activate first -cd sbd-thrust -SBD_GPU_ARCH=cc100 pip install -e . --no-build-isolation +# BLAS: defaults to whatever the linker finds, including a +# conda-installed OpenBLAS in $CONDA_PREFIX/lib. Set these to select +# a specific build (e.g. an arch-tuned OpenBLAS) +export BLAS_LIB_PATH=/path/to/blas/lib +export BLAS_LIBS=openblas # or mkl_rt ``` -Builds the CPU backend always. Adds the Thrust GPU backend when NVHPC -`nvc++` is on PATH; otherwise CPU-only. Devices: `'cpu'` and `'gpu'`. - -**Profile 2 — OpenMP target-offload GPU** (use a **separate** venv + -checkout from Profile 1) - +### Build Using the Host MPI ```bash -# one-time setup -git clone --recurse-submodules https://github.com/Qiskit/sbd-eigensolver-python.git sbd-omp-offload -python -m venv ~/venvs/sbd-omp-offload -source ~/venvs/sbd-omp-offload/bin/activate -MPICC=$(which mpicc) pip install --no-binary=mpi4py mpi4py pybind11 numpy setuptools wheel - -# build (re-run after pulling main) -source ~/venvs/sbd-omp-offload/bin/activate # always activate first -cd sbd-omp-offload -SBD_BUILD_BACKEND=gpu_omp_offload SBD_GPU_ARCH=cc100 \ - pip install -e . --no-build-isolation -``` +# Create a conda env +conda create -y -n sbd -c conda-forge \ + python=3.13.12 pybind11 numpy setuptools wheel openblas pyscf pip +# ...plus llvm-openmp on macOS -Builds the OMP-offload GPU backend only. Device: `'gpu-omp'`. +conda activate sbd # always activate first -After both profiles are installed, switching is a one-liner -(`source ~/venvs//bin/activate`) — no rebuild needed. - -**Multi-arch fat binary**: comma-separate the arches: -`SBD_GPU_ARCH=cc80,cc90,cc100`. nvc++ embeds one SASS cubin per arch -and picks the matching one at runtime. +# Install mpi4py against the host MPI +# For any GPU backend the host MPI must be GPU-aware: +# CUDA-aware on NVIDIA, ROCm-aware on AMD. +export MPI_HOME=/path/to/mpi +MPICC=$MPI_HOME/bin/mpicc python -m pip install --no-binary=mpi4py --no-cache-dir mpi4py +# NOTE: If you need only the CPU backend and no host MPI is available, let conda +# pick a compatible one with the command below (it is not GPU-aware). +# conda install -y -c conda-forge mpi4py -#### Advanced `SBD_BUILD_BACKEND` overrides +# confirm which MPI mpi4py uses -- setup.py builds against exactly this +python -c "from mpi4py import MPI; print(MPI.Get_library_version())" -Only needed when you want to deviate from the two profiles above. +# only for the SQD examples (python/examples/run_sqd_sbd.py and .ipynb) +pip install "qiskit-addon-sqd>=0.13.1" -| Value | Builds | -|---|---| -| *unset* (default) | CPU always; Thrust GPU if `nvc++` found. The "Profile 1" default. | -| `cpu` | CPU only — skip GPU even if `nvc++` is present. | -| `gpu` | Thrust GPU only — skip CPU. Errors if `nvc++` missing. | -| `both` | CPU + Thrust GPU. Errors instead of falling back if `nvc++` missing. | -| `gpu_omp_offload` | OMP-offload GPU only. The "Profile 2" install. | +# install sbd-eigensolver +pip install sbd-eigensolver -**Reverting to the LLVM/clang offload path:** prior versions of this -repo supported a separate `_core_gpu_omp_nvidia` backend built with -LLVM-with-NVPTX clang. That path was removed in v1.6 to reduce the -software prereq surface (LLVM trunk had to be source-built, NVHPC's -nvc++ does not). The tag `v1.5.0-llvm` preserves the last revision -with that backend; check it out and follow the `SETUP_LLVM_OFFLOAD.txt` -recipe there if you need the clang path back. +``` ### Verify ```bash python -c "import sbd; print(sbd.available_backends())" # CPU only: ['cpu'] -# CPU + NVHPC Thrust: ['cpu', 'gpu'] +# NVIDIA, default build: ['cpu', 'gpu', 'gpu-omp'] +# AMD, default build: ['cpu', 'gpu-omp'] # OMP-offload-only install: ['gpu-omp'] ``` -## Usage +On a GPU build, confirm which vendor and architecture the offload backend +targets — one `'gpu-omp'` device serves both vendors, so the name alone does not +say: -### Quick Start - -```python -import sbd - -# No explicit init() needed — auto-initializes on first use -config = sbd.TPB_SBD() -config.adet_comm_size = 2 -config.bdet_comm_size = 2 -config.max_it = 100 -config.eps = 1e-4 - -results = sbd.tpb_diag_from_files( - fcidumpfile='data/h2o/fcidump.txt', - adetfile='data/h2o/h2o-1em4-alpha.txt', - sbd_data=config, -) - -print(f"Energy: {results['energy']:.10f} Hartree") -sbd.finalize() +```bash +python -c "import sbd; print(sbd.get_backend('gpu-omp').__sbd_offload_target__)" +# amdgcn-amd-amdhsa:gfx90a AMD MI250X +# nvptx64-nvidia-cuda:cc90 NVIDIA H100 ``` -### Runtime backend switching - -Compatible backends coexist as separate `_core_*.so` modules and load -at `import sbd` into independent pybind11 namespaces. CPU + Thrust GPU -can co-load; the OMP-offload backend cannot (different OpenMP runtime — -see the build section). Pick one per call with the `device` parameter: +The Thrust backend is stamped too (`cuda:cc90`); the CPU backend reports `None`. -```python -import sbd - -# All compiled backends are auto-loaded -sbd.available_backends() -# CPU + Thrust install: ['cpu', 'gpu'] -# OMP-offload-only install: ['gpu-omp'] - -# Per-call override — auto-initializes on first use -result_cpu = sbd.tpb_diag(..., device='cpu') -result_thrust = sbd.tpb_diag(..., device='gpu') -result_omp = sbd.tpb_diag(..., device='gpu-omp') - -# Or set a default device explicitly (optional) -sbd.init(device='gpu') # default = NVHPC Thrust -result = sbd.tpb_diag(...) - -# Or get the backend module directly -backend = sbd.get_backend('gpu-omp') -fcidump = backend.LoadFCIDump('fcidump.txt') -``` - -In `auto` mode (the default), `_resolve_device('auto')` picks the first -compiled GPU backend in the order Thrust → OMP-offload → CPU. +## Examples -### Resource Cleanup +Located in `python/examples/`: -```python -results = sbd.tpb_diag_from_files(...) -sbd.finalize() # optional — syncs GPU and resets state -``` +- **`run_sbd_diag.py`** — Standalone TPB diagonalization (no Qiskit dependency) +- **`run_sqd_sbd.ipynb`** — Jupyter Notebook SQD loop with SBD solver (random or hardware bitstrings) +- **`run_sqd_sbd.py`** — SQD loop with SBD solver (random or hardware bitstrings) -`finalize()` calls `cudaDeviceSynchronize()` on GPU backends and resets Python state. It does **not** call `cudaDeviceReset()` (avoids CUDA-aware MPI conflicts) or `MPI_Finalize()` (handled by mpi4py). +See [python/examples/README.md](python/examples/README.md) for usage details. ## Integration with qiskit-addon-sqd @@ -309,31 +246,13 @@ result = diagonalize_fermionic_hamiltonian( See `python/examples/run_sqd_sbd.py` for a complete example. -### Comparison with qiskit-addon-dice-solver - -| Feature | dice-solver | SBD | -|---------|------------|-----| -| Solver | DICE (subprocess) | SBD (in-process) | -| GPU | No | Yes (CUDA) | -| MPI | Spawns processes | Direct integration | -| I/O | Temp files | In-memory | - -## Examples - -Located in `python/examples/`: - -- **`run_sbd_diag.py`** — Standalone TPB diagonalization (no Qiskit dependency) -- **`run_sqd_sbd.py`** — SQD loop with SBD solver (random or hardware bitstrings) - -See [python/examples/README.md](python/examples/README.md) for usage details. - ## API Reference ### Initialization | Function | Description | |----------|-------------| -| `sbd.init(device, comm_backend)` | **Optional.** Initialize MPI, set default device (`'cpu'`, `'gpu'`, `'auto'`). Auto-called on first use with defaults. | +| `sbd.init(device, comm_backend)` | **Optional.** Initialize MPI, set default device (`'cpu'`, `'gpu'`, `'gpu-omp'`, `'auto'`). Auto-called on first use with defaults. | | `sbd.finalize()` | Sync GPU, reset state. Does not call `MPI_Finalize` | | `sbd.is_initialized()` | Check init status | @@ -426,44 +345,43 @@ words in canonical order, then `n_dets` `float64` amplitudes. The optional `device` parameter overrides the default set by `init()`. -### Utilities - -```python -fcidump = sbd.LoadFCIDump("fcidump.txt", device=None) -dets = sbd.LoadAlphaDets("alphadets.txt", bit_length, total_bit_length, device=None) -string = sbd.makestring(det, bit_length, total_bit_length, device=None) -det = sbd.from_string(s, bit_length, total_bit_length, device=None) -dets = sbd.sort_bitarray(dets, device=None) -sbd.print_info() -``` - ## Backend Architecture -- Each backend is a separate pybind11 module compiled from the same `python/bindings.cpp` source with different `-D` macros (`SBD_THRUST` for the Thrust path, `USE_GPU + USE_OMP_OFFLOAD` for OMP-offload, neither for CPU). The Thrust and OMP-offload paths both compile with NVHPC `nvc++` (with `-cuda` and `-mp=gpu` respectively); CPU compiles with gcc/clang. Distinct C++ namespaces — no symbol collision when multiple coexist. +- Each backend is a separate pybind11 module compiled from the same `python/bindings.cpp` source with different `-D` macros (`SBD_THRUST` for the Thrust path, `USE_GPU + USE_OMP_OFFLOAD` for OMP-offload, neither for CPU). On NVIDIA the Thrust and OMP-offload paths both compile with NVHPC `nvc++` (`-cuda` and `-mp=gpu` respectively); on AMD the OMP-offload path compiles with ROCm `amdclang++` (`-fopenmp-targets=amdgcn-amd-amdhsa --offload-arch=gfx*`); CPU compiles with gcc/clang. Distinct C++ namespaces — no symbol collision when multiple coexist. +- **`'gpu-omp'` is vendor-neutral by design:** one module, one device string, for both NVIDIA and AMD. It is the same source with the same macros and only the compiler differs, and since no wheels are published — every install compiles on the target machine — an install serves one GPU vendor. `__sbd_offload_target__` records which one. Aliases (`gpu-amd-omp`, `gpu-rocm-omp`, `rocm`, `gpu-nvidia-omp`, …) resolve to it so a vendor-flavoured guess lands correctly. +- **Thrust is NVIDIA-only.** Upstream SBD wires that path to `nvc++ -cuda`, so there is no rocThrust configuration to build; `SBD_BUILD_BACKEND=gpu` fails fast on AMD rather than quietly producing a CPU-only install under a name that says `gpu`. - `get_backend(device)` resolves the `device=` string and returns the appropriate module; all wrapper functions accept an optional `device` parameter. Aliases for back-compat live in `sbd._device_aliases`. - GPU device assignment: `gpu_id = mpi_rank % num_gpus` (set per `tpb_diag()` call in `bindings.cpp`); same logic for both Thrust and OMP-offload paths. - Backends differ in which phases run on the GPU vs the host. Davidson and the matvec (`mult`) live on the GPU under both Thrust and OMP-offload. The diagonal-Hamiltonian preconditioner (`makeQChamDiagTerms`) is GPU-resident under Thrust but runs on the host under OMP-offload (no `#pragma omp target` port in `tpb/qcham.h`). +- Verify what GPU architecture is supported in the binary: + ``` + NVIDIA: cuobjdump --list-elf + AMD: llvm-objdump --offloading + ``` + Or just ask the module what it was built for: + ``` + python -c "import sbd; \ + print(sbd.get_backend('gpu-omp').__sbd_offload_target__)" + -> amdgcn-amd-amdhsa:gfx90a + ``` ## Troubleshooting -**`ImportError` on macOS (symbol not found):** Python's libc++ and Homebrew clang's libc++ may differ. Use system clang: `CC=/usr/bin/clang CXX=/usr/bin/clang++`. +**GPU backends silently build as host code:** a conda compiler package +(`cxx-compiler`, `gxx_linux-64`, `clangxx_osx-*`) sets `CC`/`CXX` on activation and +the build respects a caller-set compiler, so `nvc++`/`amdclang++` never run. Unset +`CC`/`CXX`, or keep conda compilers out of the build env. -**`ImportError: _core_cpu`:** Backend not built. Rebuild: `pip install -e . --no-build-isolation -v` - -**GPU not building:** Check `which nvc++` and set `NVHPC_HOME`. +**GPU not building:** On NVIDIA check `which nvc++` and set `NVHPC_HOME`. On AMD +check `which amdclang++` and set `ROCM_HOME`. The build prints which toolchain it +picked (`Found amdclang++ in PATH: …` / `Found NVIDIA HPC SDK at: …`) and, for +the offload backend, the resolved architecture; on a host with both toolchains +force the choice with `SBD_GPU_VENDOR=amd|nvidia`. **MPI errors:** Verify `MPI_HOME`, check `python -c "from mpi4py import MPI; print(MPI.Get_version())"`. -**OMP-offload runs all land on GPU 0 in multi-GPU jobs:** symptom — every MPI rank shows large memory only on GPU 0 in `nvidia-smi`. The bindings call `omp_set_default_device(mpi_rank % n_dev)`, but `omp_get_num_devices()` can return 0 in some dlopen scenarios. The bindings fall back to parsing `CUDA_VISIBLE_DEVICES` to recover the device count, so make sure that env var is exported and lists all your GPUs (e.g. `0,1,2,3`). Slurm/`srun --gres=gpu:N` and OpenMPI's default binding policy already do this; if you've custom-restricted `CUDA_VISIBLE_DEVICES` to a single GPU per rank, set it manually before launch. - -**OMP-offload + UCX MPI fails with `MPI_INIT failed`:** mpi4py 4.x requests `MPI_THREAD_MULTIPLE` by default, which UCX in HPCX rejects with `UCP worker does not support MPI_THREAD_MULTIPLE`. Set `MPI4PY_RC_THREAD_LEVEL=serialized` (or `funneled`/`single`) in the environment, or `import mpi4py; mpi4py.rc.thread_level = 'serialized'` before `from mpi4py import MPI`. - -## Performance Tips - -**CPU:** `OMP_NUM_THREADS` = cores per MPI rank. -**GPU (Thrust):** 1 rank per GPU, `OMP_NUM_THREADS=1`, use method 0 (matrix-free Davidson). -**GPU (OMP-offload):** 1 rank per GPU, `OMP_NUM_THREADS` ≈ socket-local cores per rank, **pin each rank to one socket** (e.g. `mpirun --map-by ppr:N:socket --bind-to socket …` or wrap with `numactl --cpunodebind=… --membind=…`). Without pinning, the host-side `makeQChamDiagTerms` loop and the host-side orchestration inside Davidson degrade ~7× and 2–3× respectively because OMP threads thrash across NUMA nodes. +**OMP-offload runs all land on GPU 0 in multi-GPU jobs:** symptom — every MPI rank shows large memory only on GPU 0 in `nvidia-smi` (or `rocm-smi`). The bindings call `omp_set_default_device(mpi_rank % n_dev)`, but `omp_get_num_devices()` can return 0 in some dlopen scenarios. The bindings fall back to counting the entries in the vendor's device-visibility variable — `CUDA_VISIBLE_DEVICES` on NVIDIA, `ROCR_VISIBLE_DEVICES` or `HIP_VISIBLE_DEVICES` on AMD — so make sure the relevant one is exported and lists all your GPUs (e.g. `0,1,2,3`). Slurm/`srun --gres=gpu:N` and OpenMPI's default binding policy already do this; if you've custom-restricted it to a single GPU per rank, set it manually before launch. ---- +**Ranks die with `Bus error` or `SIGSEGV` inside the MPI's own copy path** (`MPIR_Localcopy`, `ucp_worker_progress`, ...) **on a GPU backend:** the MPI is not GPU-aware and was handed a device pointer. Rebuild UCX `--with-cuda` / `--with-rocm`, and confirm with `ucx_info -d | grep -i 'Transport: cuda'` (or `rocm`). Two things mislead here. A partly GPU-aware stack fails in only one place: an MPICH with GPU support *disabled* over a CUDA-aware UCX ran OMP-offload fine and crashed only in Thrust, because the inter-rank path went through UCX while the local-copy path did not. And on AMD a non-ROCm-aware MPI does not crash at all — ROCm maps device memory into the process address space, so the host copy succeeds and merely stages everything through the host aperture (measured on MI250X, XNACK off, 8 ranks) — so a working AMD run is not evidence that the MPI is ROCm-aware. **Repository:** https://github.com/Qiskit/sbd-eigensolver-python diff --git a/python/__init__.py b/python/__init__.py index 824f0a4..81e0331 100644 --- a/python/__init__.py +++ b/python/__init__.py @@ -28,8 +28,9 @@ """ import os -import subprocess +import re from importlib.metadata import version +import subprocess # Single-sourced from the [project] table in pyproject.toml, so the version is # declared in exactly one place. Note the argument is the *distribution* name @@ -37,34 +38,207 @@ __version__ = version("sbd-eigensolver") # --------------------------------------------------------------------------- -# Backend registry — eagerly load all available backends at import time. -# All compiled backends can coexist: separate .so files with separate -# pybind11 namespaces, no global C++ state conflicts. +# Backend registry. +# +# Backends are resolved lazily: only the one actually requested gets imported. +# This is not a micro-optimisation. Every _core_*.so here links NVHPC's +# libnvomp, and _core_cpu is built without -mp=gpu; loading it first leaves +# that runtime initialised host-only, after which the OMP-offload backend +# cannot acquire a device. It does not fail -- offload regions quietly run on +# the host while device queries still report a GPU, so the run looks +# accelerated, returns the correct energy, and exits 0. Measured on GB200: +# with _core_cpu co-loaded, OMP_TARGET_OFFLOAD=MANDATORY aborts with "Could +# not run target region"; importing only _core_gpu_omp_offload from the very +# same directory succeeds on the device. Loading one backend per process is +# therefore what makes co-resident .so files safe, and lets a single +# environment serve all three. +# +# The AMD build has the same SHAPE of hazard -- there _core_cpu and +# _core_gpu_omp_offload share LLVM's libomp, with only the offload one pulling +# in libomptarget -- so the same one-backend-per-process discipline applies. It +# has not been reproduced on AMD, and it should not need to be: the discipline +# below prevents it either way, and OMP_TARGET_OFFLOAD=MANDATORY (set before the +# import, further down) turns a demotion to the host into an abort rather than a +# plausible wrong answer. # --------------------------------------------------------------------------- -_backends = {} -# Device-string aliases. Keys are user-facing strings; values are the -# canonical key in _backends. e.g. 'gpu-omp' resolves to whatever -# OMP-offload backend is built. -_device_aliases = {} +# (module name, canonical device, aliases) +# +# 'gpu-omp' is vendor-NEUTRAL: the same module serves NVIDIA (nvc++ -mp=gpu) and +# AMD (amdclang++ --offload-arch=gfx*), since it is the same source with the same +# macros and only the compiler differs. The vendor-flavoured names are aliases so +# that guessing 'gpu-amd-omp' or 'rocm' lands on the right backend instead of +# raising; get_backend('gpu-omp').__sbd_offload_target__ reports what a given +# install was actually built for. +_BACKEND_SPECS = ( + ('_core_cpu', 'cpu', ()), + ('_core_gpu_thrust', 'gpu', ('gpu-thrust', 'gpu-nvidia', 'cuda')), + ('_core_gpu_omp_offload', 'gpu-omp', ('gpu-omp-offload', 'gpu-nvhpc-omp', + 'gpu-nvidia-omp', 'gpu-amd-omp', + 'gpu-rocm-omp', 'rocm')), +) + +_backends = {} # device -> imported module, populated on first use +_backend_errors = {} # device -> why it is unusable +_device_aliases = {a: dev for _m, dev, al in _BACKEND_SPECS for a in al} +_usable_cache = None # list of usable devices; the static scan runs once + + +def _extension_path(module_name): + """Path of the built extension for `module_name`, or None.""" + import os + here = os.path.dirname(os.path.abspath(__file__)) + try: + names = os.listdir(here) + except OSError: + return None + for n in sorted(names): + if n.startswith(module_name + '.') and n.endswith(('.so', '.pyd', '.dylib')): + return os.path.join(here, n) + return None + + +def _inspect_extension(path): + """Statically judge whether `path` can load. Returns (ok, reason). + + Deliberately does NOT import it. Importing a backend runs MPI_Init, which + on an MPI without PMIx support hangs when the process was not started under + a launcher -- a diagnostic that hangs is worse than none. `ldd` inspects the + file without executing its initialisers. + + Catches a missing dependency and an architecture mismatch. It cannot catch + an ABI/symbol mismatch, which only surfaces on a real dlopen; that is + reported by _load_backend() at first use. + """ + import subprocess, sys as _sys + if _sys.platform == 'darwin': + return True, None # otool output does not mark unresolved deps + try: + out = subprocess.run(['ldd', path], capture_output=True, text=True, + timeout=20).stdout + except Exception: + return True, None # no ldd: fall through to the real import + if 'not a dynamic executable' in out: + try: + kind = subprocess.run(['file', '-b', path], capture_output=True, + text=True, timeout=20).stdout.strip() + except Exception: + kind = 'unknown' + import platform + return False, (f"architecture mismatch: the file is '{kind}' but this " + f"host is {platform.machine()}. Rebuild it here. (Note " + f"the loader reports this as 'cannot open shared object " + f"file', which is misleading -- the file exists.)") + missing = [l.split('=>')[0].strip() for l in out.splitlines() if 'not found' in l] + if missing: + return False, (f"missing shared librar{'y' if len(missing) == 1 else 'ies'}: " + f"{', '.join(missing)}. Put the containing directory on " + f"LD_LIBRARY_PATH, or rebuild so RPATH covers it.") + return True, None + + +def _scan_backends(): + """Populate the usable-device list and the reasons for the rest. + + Runs once. A backend that was simply not built is recorded but not warned + about -- building CPU-only is a normal choice. A backend that IS built but + cannot load is a broken install, so that one warns. + """ + global _usable_cache + if _usable_cache is not None: + return _usable_cache + import warnings + usable, broken = [], [] + for module_name, device, _aliases in _BACKEND_SPECS: + path = _extension_path(module_name) + if path is None: + _backend_errors[device] = ( + f"not built (no {module_name}.*.so in the package directory)" + ) + continue + ok, reason = _inspect_extension(path) + if ok: + usable.append(device) + else: + _backend_errors[device] = reason + broken.append(f"{device}: {reason}") + _usable_cache = usable + if broken: + warnings.warn( + "sbd: backend(s) present on disk but not usable —\n " + + "\n ".join(broken), + RuntimeWarning, + stacklevel=3, + ) + return usable + +def _load_backend(device): + """Import the backend for `device`, or raise with an actionable message.""" + if device in _backends: + return _backends[device] + module_name = next((m for m, d, _a in _BACKEND_SPECS if d == device), None) + if module_name is None: + raise RuntimeError(f"Unknown backend device '{device}'.") + + path = _extension_path(module_name) + if path is None: + raise RuntimeError( + f"sbd backend '{device}' is not built.\n" + f" expected: {module_name}.*.so in the sbd package directory\n" + f" usable : {_scan_backends() or 'none'}\n" + f" fix : rebuild, e.g. " + f"SBD_BUILD_BACKEND={'cpu' if device == 'cpu' else 'gpu_omp_offload' if device == 'gpu-omp' else 'gpu'}" + f" pip install -e . --no-build-isolation" + ) -def _try_load(module_name, primary_device, *aliases): + if device == 'gpu-omp': + # Make a host fallback fail loudly instead of succeeding quietly. + # + # OMP_TARGET_OFFLOAD=DEFAULT -- the OpenMP default -- means "device if + # one is available, else host". A demoted run returns the correct energy + # and exits 0, with device queries still reporting a GPU, so it looks + # accelerated. Reachable whenever no device is available to the rank: an + # empty CUDA_VISIBLE_DEVICES, a mispinning launcher wrapper, a GPU-less + # node. MANDATORY turns each of those into an immediate error. + # + # This is a standard OpenMP 5.0 ICV, honoured by NVHPC libnvomp, GNU + # libgomp and LLVM libomptarget alike -- not an NVHPC extension. It must + # be set before the runtime reads it, which is why it happens here, + # immediately ahead of the import. + # + # Only when the caller has not chosen: an explicit + # OMP_TARGET_OFFLOAD=DEFAULT still wins, for anyone who wants host + # fallback (a smoke test on a GPU-less login node). An empty value is + # treated as unset -- `export OMP_TARGET_OFFLOAD=` is not a request for + # opportunistic offload, and setdefault() would otherwise honour it. + import os as _os + if not _os.environ.get('OMP_TARGET_OFFLOAD'): + _os.environ['OMP_TARGET_OFFLOAD'] = 'MANDATORY' + + from importlib import import_module try: - from importlib import import_module mod = import_module(f'.{module_name}', package=__name__) - except ImportError: - return False - _backends[primary_device] = mod - for a in aliases: - _device_aliases[a] = primary_device - return True + except Exception as exc: + # Never substitute a different backend: a silently-swapped device is the + # same class of lie as a silent host fallback. + ok, reason = _inspect_extension(path) + detail = reason or f"{type(exc).__name__}: {exc}" + _backend_errors[device] = detail + raise RuntimeError( + f"sbd backend '{device}' could not be loaded.\n" + f" module: {path}\n" + f" cause : {type(exc).__name__}: {exc}\n" + f" detail: {detail}\n" + f" If the cause mentions an undefined symbol, the extension was " + f"built against a different Python or mpi4py than this " + f"environment's; rebuild it here." + ) from exc + _backends[device] = mod + _backend_errors.pop(device, None) + return mod -_try_load('_core_cpu', 'cpu') -_try_load('_core_gpu_thrust', 'gpu', 'gpu-thrust', 'gpu-nvidia', 'cuda') -_try_load('_core_gpu_omp_offload', 'gpu-omp', 'gpu-omp-offload', - 'gpu-nvhpc-omp', 'gpu-nvidia-omp') # --------------------------------------------------------------------------- # Global session state @@ -80,17 +254,38 @@ def _try_load(module_name, primary_device, *aliases): def _gpu_available(): - """Check if GPU is available via nvidia-smi (cached).""" + """Check whether any GPU is present, via nvidia-smi or rocm-smi (cached). + + rocm-smi is consulted as well as nvidia-smi because the 'gpu-omp' backend + serves AMD too. Checking only nvidia-smi made device='auto' resolve to 'cpu' + on an AMD host that had a perfectly good OMP-offload backend built -- a + silent downgrade to the slow path. + + This asks about HARDWARE only. Whether a backend was compiled for it is a + separate question, answered by _scan_backends(); _resolve_device() needs both. + """ global _gpu_check_cache if _gpu_check_cache is not None: return _gpu_check_cache - try: - result = subprocess.run( - ['nvidia-smi'], capture_output=True, timeout=2 - ) - _gpu_check_cache = result.returncode == 0 - except Exception: - _gpu_check_cache = False + _gpu_check_cache = False + # Per-probe timeouts so nvidia-smi cannot inherit rocm-smi's 30 s, and + # rocm-smi needs a GPU[] line since it can exit 0 with no GPU present. + for probe, args, timeout in ( + ('nvidia-smi', [], 5), + ('rocm-smi', ['--showid'], 30), + ): + try: + result = subprocess.run( + [probe, *args], capture_output=True, text=True, timeout=timeout + ) + if result.returncode != 0: + continue + if probe == 'rocm-smi' and not re.search(r'GPU\[\d+\]', result.stdout or ''): + continue + _gpu_check_cache = True + break + except Exception: + continue return _gpu_check_cache @@ -99,15 +294,15 @@ def _resolve_device(device): Auto-resolution prefers Thrust GPU (`'gpu'`) over OMP-offload (`'gpu-omp'`) when both are built and a GPU is present, since the - Thrust path is the long-validated default. In practice OMP-offload - is built into its own venv/install (different OpenMP runtime, can't - co-load with Thrust/CPU), so this branch only fires when an OMP-only - install is in use. + Thrust path is the long-validated default. Since backends load lazily, + all three may live in one install: only the resolved one is imported, so + _core_cpu never poisons the OMP-offload runtime. """ if device == 'auto': - if 'gpu' in _backends and _gpu_available(): + usable = _scan_backends() + if 'gpu' in usable and _gpu_available(): return 'gpu' - if 'gpu-omp' in _backends and _gpu_available(): + if 'gpu-omp' in usable and _gpu_available(): return 'gpu-omp' return 'cpu' return _device_aliases.get(device, device) @@ -126,8 +321,11 @@ def init(device='cpu', comm_backend='mpi'): Args: device: Default compute device — 'cpu', 'gpu', 'gpu-omp', or 'auto'. + 'gpu' is the NVIDIA-only Thrust backend; 'gpu-omp' is OpenMP + target offload and serves NVIDIA and AMD alike. Aliases: 'gpu-thrust' / 'gpu-nvidia' / 'cuda' (= 'gpu'); - 'gpu-omp-offload' / 'gpu-nvhpc-omp' / 'gpu-nvidia-omp' (= 'gpu-omp'). + 'gpu-omp-offload' / 'gpu-nvhpc-omp' / 'gpu-nvidia-omp' / + 'gpu-amd-omp' / 'gpu-rocm-omp' / 'rocm' (= 'gpu-omp'). comm_backend: Communication backend — 'mpi'. Raises: @@ -138,7 +336,7 @@ def init(device='cpu', comm_backend='mpi'): if _initialized: return # already initialized — silently no-op - if not _backends: + if not _scan_backends(): raise RuntimeError( "No SBD backends available. Build with:\n" " pip install -e . --no-build-isolation (auto: CPU + Thrust GPU)\n" @@ -162,11 +360,12 @@ def init(device='cpu', comm_backend='mpi'): # Resolve default device resolved = _resolve_device(device) - if resolved not in _backends: - available = list(_backends.keys()) + usable = _scan_backends() + if resolved not in usable: + why = _backend_errors.get(resolved, 'unknown reason') raise RuntimeError( - f"Device '{resolved}' requested but backend not available. " - f"Available: {available}" + f"Device '{resolved}' requested but its backend is not usable " + f"({why}). Usable: {usable or 'none'}" ) _default_device = resolved _initialized = True @@ -225,12 +424,7 @@ def get_backend(device=None): if device is None: device = _default_device or 'auto' device = _resolve_device(device) - if device not in _backends: - available = list(_backends.keys()) - raise RuntimeError( - f"Backend '{device}' not available. Available: {available}" - ) - return _backends[device] + return _load_backend(device) def _ensure_initialized(): @@ -426,10 +620,51 @@ def gdb_diag(fcidump, det, sbd_data, # --------------------------------------------------------------------------- def available_backends(): - """Get list of compiled backends ('cpu', 'gpu', 'gpu-omp').""" + """Devices whose extension is built and structurally loadable. + + Each candidate .so is inspected statically -- present on disk, no + unresolved shared-library dependencies, matching architecture -- without + importing it. Importing would run MPI_Init, which hangs on an MPI lacking + PMIx support when the process was not started under a launcher. + + A listed backend is therefore "present and structurally sound", not + "guaranteed to load": an ABI/symbol mismatch only shows up on a real + dlopen, and surfaces from get_backend(). Reasons for anything excluded are + in backend_load_errors(); a backend that is merely not built is recorded + there but does not warn, since building a subset is normal. + """ + return list(_scan_backends()) + + +def loaded_backends(): + """Devices actually imported into this process so far. + + Normally one: backends load on first use. More than one means something + imported them explicitly, which is worth knowing -- co-loading _core_cpu + with the OMP-offload backend silently demotes offload to the host. + """ return list(_backends.keys()) +def backend_load_errors(): + """Why each unusable backend is unusable. + + Maps device -> reason, distinguishing "not built" from a missing shared + library, an architecture mismatch, or a failed import. + """ + return dict(_backend_errors) + + +def has_backend_conflict(): + """True when _core_cpu and the OMP-offload backend are both loaded here. + + That combination leaves libnvomp initialised host-only, so offload regions + run on the host while device queries still report a GPU. Lazy loading + normally prevents it; this catches a caller that imported both explicitly. + """ + return 'cpu' in _backends and 'gpu-omp' in _backends + + def print_info(): """Print SBD information.""" print("=" * 60) @@ -465,6 +700,10 @@ def print_info(): # Backend access 'get_backend', + 'available_backends', + 'loaded_backends', + 'backend_load_errors', + 'has_backend_conflict', # Query 'get_device', diff --git a/python/bindings.cpp b/python/bindings.cpp index 5d4bb13..0024ccc 100644 --- a/python/bindings.cpp +++ b/python/bindings.cpp @@ -14,10 +14,16 @@ * @file python/bindings.cpp * @brief Python bindings for SBD TPB diagonalization using pybind11 * - * This file is compiled three times with different module names + flags: + * This file is compiled once per backend, with different module names + flags: * - _core_cpu : CPU backend (host OpenMP via -fopenmp) * - _core_gpu_thrust : Thrust GPU backend (with -DSBD_THRUST, nvc++ -cuda) - * - _core_gpu_omp_offload : OpenMP-offload GPU (with -DUSE_OMP_OFFLOAD, nvc++ -mp=gpu) + * - _core_gpu_omp_offload : OpenMP-offload GPU (with -DUSE_OMP_OFFLOAD), built by + * nvc++ -mp=gpu on NVIDIA, or amdclang++ + * --offload-arch=gfx* on AMD. ONE module and one + * 'gpu-omp' device string serve both vendors -- + * same source, same macros, different compiler -- + * with the target recorded as + * __sbd_offload_target__ (see SBD_OFFLOAD_TARGET). * * The module name is controlled by the SBD_MODULE_NAME macro. */ @@ -54,6 +60,55 @@ MPI_Comm get_mpi_comm(py::object py_comm) { return *comm_ptr; } +#ifdef USE_OMP_OFFLOAD +/** + * Pin this rank to one offload device: device = mpi_rank % n_devices. + * + * Note: when this .so is loaded via Python dlopen, the symbol + * omp_get_num_devices binds to libomp.so's stub (which returns 0 because libomp + * itself doesn't manage offload devices) instead of libomptarget's working + * version. omp_set_default_device IS shared between the two, so once we know the + * count we can still set the device correctly. So fall back to counting the + * entries in the vendor's device-visibility variable when the count reads 0. + * + * The variable to read is vendor-specific, and the order is decided at COMPILE + * time from the target this module was built for rather than by probing all of + * them: that keeps the NVIDIA build reading exactly CUDA_VISIBLE_DEVICES first, + * as it always has. The other vendor's names are still listed as a fallback, + * which costs nothing (they are unset on a single-vendor host) and helps on an + * oddly-configured node. + */ +static void sbd_pin_offload_device(int mpi_rank) { + int n_dev = omp_get_num_devices(); + if (n_dev <= 0) { + static const char* const kVisibleVars[] = { +#ifdef SBD_OFFLOAD_VENDOR_AMD + "ROCR_VISIBLE_DEVICES", // AMD: honoured by the ROCm OMP runtime + "HIP_VISIBLE_DEVICES", // AMD: HIP-level equivalent + "CUDA_VISIBLE_DEVICES", +#else + "CUDA_VISIBLE_DEVICES", // NVIDIA + "ROCR_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", +#endif + }; + for (const char* var : kVisibleVars) { + const char* val = std::getenv(var); + if (val && *val) { + n_dev = 1; + for (const char* p = val; *p; ++p) { + if (*p == ',') ++n_dev; + } + break; + } + } + } + if (n_dev > 0) { + omp_set_default_device(mpi_rank % n_dev); + } +} +#endif + // Module name is set by compiler flag, e.g. // -DSBD_MODULE_NAME=_core_cpu | _core_gpu_thrust | _core_gpu_omp_offload #ifndef SBD_MODULE_NAME @@ -68,6 +123,21 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { m.doc() = "Python bindings for SBD (Selected Basis Diagonalization) library - CPU backend"; #endif + // Which GPU target this module was compiled for: + // gpu-omp on AMD "amdgcn-amd-amdhsa:gfx90a" + // gpu-omp on NVIDIA "nvptx64-nvidia-cuda:cc90" + // gpu (Thrust) "cuda:cc90" + // cpu None + // For the OMP-offload backend this is the only way to tell the vendor apart, + // since one module and one device string ('gpu-omp') serve both. For Thrust + // the vendor is never in doubt but the architecture is -- and an arch + // mismatch is the usual reason a module built elsewhere will not run here. +#ifdef SBD_OFFLOAD_TARGET + m.attr("__sbd_offload_target__") = py::str(SBD_OFFLOAD_TARGET); +#else + m.attr("__sbd_offload_target__") = py::none(); +#endif + // Initialize mpi4py if (import_mpi4py() < 0) { throw std::runtime_error("Failed to import mpi4py"); @@ -259,29 +329,7 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { #endif #ifdef USE_OMP_OFFLOAD // Assign OMP-offload device based on MPI rank. - // - // Note: when this .so is loaded via Python dlopen, the symbol - // omp_get_num_devices binds to libomp.so's stub (which returns 0 - // because libomp itself doesn't manage offload devices) instead - // of libomptarget's working version. omp_set_default_device - // IS shared between the two, so once we know the count we can - // still set the device correctly. Fall back to parsing - // CUDA_VISIBLE_DEVICES when omp_get_num_devices reports 0. - { - int n_dev = omp_get_num_devices(); - if (n_dev <= 0) { - const char* cvd = std::getenv("CUDA_VISIBLE_DEVICES"); - if (cvd && *cvd) { - n_dev = 1; - for (const char* p = cvd; *p; ++p) { - if (*p == ',') ++n_dev; - } - } - } - if (n_dev > 0) { - omp_set_default_device(mpi_rank % n_dev); - } - } + sbd_pin_offload_device(mpi_rank); #endif // Output variables. Since upstream PR#71 the TPB det lists are @@ -381,6 +429,10 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { myDevice = mpi_rank % numDevices; hipSetDevice(myDevice); #endif +#endif +#ifdef USE_OMP_OFFLOAD + // Assign OMP-offload device based on MPI rank. + sbd_pin_offload_device(mpi_rank); #endif // Output variables @@ -488,29 +540,7 @@ PYBIND11_MODULE(SBD_MODULE_NAME, m) { #endif #ifdef USE_OMP_OFFLOAD // Assign OMP-offload device based on MPI rank. - // - // Note: when this .so is loaded via Python dlopen, the symbol - // omp_get_num_devices binds to libomp.so's stub (which returns 0 - // because libomp itself doesn't manage offload devices) instead - // of libomptarget's working version. omp_set_default_device - // IS shared between the two, so once we know the count we can - // still set the device correctly. Fall back to parsing - // CUDA_VISIBLE_DEVICES when omp_get_num_devices reports 0. - { - int n_dev = omp_get_num_devices(); - if (n_dev <= 0) { - const char* cvd = std::getenv("CUDA_VISIBLE_DEVICES"); - if (cvd && *cvd) { - n_dev = 1; - for (const char* p = cvd; *p; ++p) { - if (*p == ',') ++n_dev; - } - } - } - if (n_dev > 0) { - omp_set_default_device(mpi_rank % n_dev); - } - } + sbd_pin_offload_device(mpi_rank); #endif // Output variables. co_adet/co_bdet are det_vector<...half> since diff --git a/python/device_config.py b/python/device_config.py index f5f9789..3b933de 100644 --- a/python/device_config.py +++ b/python/device_config.py @@ -17,6 +17,7 @@ without changing user code. """ +import re import subprocess @@ -47,9 +48,10 @@ def __init__(self, device: str = 'cpu', Initialize device configuration. Args: - device: Backend device key — 'cpu', 'gpu' (NVHPC Thrust), - 'gpu-omp' (nvc++ OpenMP target offload), or any alias - known to ``sbd._device_aliases``. Default 'cpu'. + device: Backend device key — 'cpu', 'gpu' (NVHPC Thrust, + NVIDIA-only), 'gpu-omp' (OpenMP target offload, NVIDIA or + AMD), or any alias known to ``sbd._device_aliases``. + Default 'cpu'. use_precalculated_dets: Use precalculated determinants (GPU only) max_memory_gb: Maximum GPU memory in GB (-1 = auto) use_gpu: Deprecated boolean. If supplied without ``device``, @@ -67,26 +69,51 @@ def __init__(self, device: str = 'cpu', @classmethod def auto(cls, max_memory_gb: int = -1) -> 'DeviceConfig': """ - Auto-detect GPU availability and use it if available. - + Auto-detect the best available backend and use it. + + Resolves against the backends that were actually COMPILED, not just the + hardware that is present. Detecting a GPU and returning 'gpu' + unconditionally was wrong in two ways: on an AMD host it selected the + Thrust/CUDA backend, which cannot exist there (upstream wires Thrust to + nvc++ -cuda), so a machine that reported "GPU detected (HIP)" then failed + to load a CUDA module; and on a CPU-only build with a GPU present it + picked a backend that was never built. + + Preference order matches sbd._resolve_device(): Thrust ('gpu') first + where it exists, since it is the long-validated NVIDIA default and keeps + more phases on the device, then OpenMP offload ('gpu-omp'), then CPU. + Args: max_memory_gb: Maximum GPU memory in GB (-1 = auto) - + Returns: - DeviceConfig configured for GPU if available, CPU otherwise + DeviceConfig for the best backend that is both built and runnable """ - # Check if CUDA or HIP is available has_cuda = cls._check_cuda() has_hip = cls._check_hip() - - use_gpu = has_cuda or has_hip - if use_gpu: - print(f"GPU detected ({'CUDA' if has_cuda else 'HIP'}), using GPU acceleration") + try: + from . import available_backends + built = available_backends() + except Exception: + built = [] + + device = 'cpu' + if has_cuda or has_hip: + vendor = 'CUDA' if has_cuda else 'HIP/ROCm' + for candidate in ('gpu', 'gpu-omp'): + if candidate in built: + device = candidate + break + if device == 'cpu': + print(f"GPU detected ({vendor}) but no GPU backend is built " + f"(available: {built or 'none'}), using CPU") + else: + print(f"GPU detected ({vendor}), using GPU acceleration " + f"via device={device!r}") else: print("No GPU detected, using CPU") - device = 'gpu' if use_gpu else 'cpu' return cls(device=device, max_memory_gb=max_memory_gb) @classmethod @@ -97,10 +124,14 @@ def cpu(cls) -> 'DeviceConfig': @classmethod def gpu(cls, use_precalculated_dets: bool = True, max_memory_gb: int = -1) -> 'DeviceConfig': - """Force NVHPC Thrust GPU execution. + """Force NVHPC Thrust GPU execution. **NVIDIA only.** + + Requires SBD compiled with THRUST (the ``_core_gpu_thrust`` extension, + i.e. ``SBD_BUILD_BACKEND=gpu``, or the default ``auto``). - Requires SBD compiled with THRUST (the ``_core_gpu`` extension, - i.e. ``SBD_BUILD_BACKEND=gpu`` or ``=both``). + There is no AMD equivalent: upstream SBD wires the Thrust path to + ``nvc++ -cuda``, so no rocThrust configuration exists to build. On an AMD + host use :meth:`gpu_omp` instead. """ return cls(device='gpu', use_precalculated_dets=use_precalculated_dets, @@ -108,14 +139,26 @@ def gpu(cls, use_precalculated_dets: bool = True, @classmethod def gpu_omp(cls, max_memory_gb: int = -1) -> 'DeviceConfig': - """Force OpenMP target-offload GPU execution. + """Force OpenMP target-offload GPU execution. Works on NVIDIA **and AMD**. Requires SBD compiled with the OMP-offload backend (the - ``_core_gpu_omp_offload`` extension, i.e. - ``SBD_BUILD_BACKEND=gpu_omp_offload``). The backend uses - ``nvc++ -mp=gpu`` with NVHPC's ``libnvomp`` runtime; it cannot - coexist in a single Python process with the CPU or Thrust GPU - backends (different OpenMP runtimes — install separately). + ``_core_gpu_omp_offload`` extension), which the default + ``SBD_BUILD_BACKEND=auto`` builds whenever a GPU compiler is present; + narrow it to ``gpu_omp_offload`` to build only this one. + + One module and one device string serve both vendors -- the same source + and macros, compiled by ``nvc++ -mp=gpu`` with NVHPC's ``libnvomp``, or by + ``amdclang++ --offload-arch=gfx*`` with LLVM's ``libomp``/``libomptarget``. + ``sbd.get_backend('gpu-omp').__sbd_offload_target__`` reports which, e.g. + ``'amdgcn-amd-amdhsa:gfx90a'``. + + It installs alongside the CPU backend (and Thrust, on NVIDIA) -- backends + are imported lazily, one per process, which is what keeps them apart. The + one combination to avoid in a single process is this backend together + with the CPU one: they share an OpenMP runtime (``libnvomp`` on NVIDIA, + ``libomp`` on AMD), and loading ``_core_cpu`` first leaves it initialised + host-only, after which offload regions run on the host. See + :func:`sbd.has_backend_conflict`. """ return cls(device='gpu-omp', max_memory_gb=max_memory_gb) @@ -156,14 +199,29 @@ def _check_cuda(cls) -> bool: @classmethod def _check_hip(cls) -> bool: - """Check if HIP/ROCm is available (cached).""" + """Check if HIP/ROCm is available (cached). + + The timeout is much longer than the CUDA probe's on purpose: rocm-smi is + a Python program that enumerates devices, measured at 1.1-1.3 s on an + 8-GCD MI250X node, where nvidia-smi answers in tens of milliseconds. The + 2 s used here originally left barely 1.5x of margin and DID flake -- + reporting no AMD GPU on a machine with eight, which then sent + DeviceConfig.auto() to the CPU backend. The result is cached, so a + generous timeout costs at most one slow call per process. + """ if cls._hip_cache is not None: return cls._hip_cache try: result = subprocess.run( - ['rocm-smi'], capture_output=True, timeout=2 + ['rocm-smi', '--showid'], capture_output=True, text=True, + timeout=30, + ) + # Require a GPU[] line, not just exit 0: rocm-smi can succeed on + # a GPU-less build host, and auto() would pick an unrunnable backend. + cls._hip_cache = ( + result.returncode == 0 + and re.search(r'GPU\[\d+\]', result.stdout or '') is not None ) - cls._hip_cache = result.returncode == 0 except Exception: cls._hip_cache = False return cls._hip_cache @@ -232,11 +290,15 @@ def get_device_info() -> dict: try: result = subprocess.run( ['rocm-smi', '--showid'], - capture_output=True, text=True, timeout=2, + capture_output=True, text=True, timeout=10, ) if result.returncode == 0: - info['gpu_count'] = len([l for l in result.stdout.split('\n') - if 'GPU' in l]) + # Count DISTINCT GPU indices, not lines mentioning "GPU": + # --showid prints several lines per device (Device Name, Device + # ID, Rev, Subsystem ID, GUID), so a line count reported 40 for + # the 8 GCDs of a 4-card MI250X node. + ids = re.findall(r'GPU\[(\d+)\]', result.stdout) + info['gpu_count'] = len(set(ids)) except Exception: pass @@ -244,19 +306,41 @@ def get_device_info() -> dict: def print_device_info(): - """Print information about available compute devices.""" + """Print the compiled backends, then the hardware they could run on. + + Backends come first because they are what actually constrains a run: the + hardware being present says nothing about whether a backend was compiled + for it. An earlier version printed "CPU Available: Always" unconditionally + while sbd.available_backends() returned [] -- reported as issue #9. + """ info = get_device_info() - - print("="*60) + + from . import available_backends, backend_load_errors, has_backend_conflict + backends = available_backends() + errors = backend_load_errors() + + print("=" * 60) print("SBD Device Information") - print("="*60) - + print("=" * 60) + + print(f"Compiled backends: {', '.join(backends) if backends else 'NONE'}") + if not backends: + print(" Nothing was built, or nothing could be loaded. This install") + print(' cannot run: solve_sci will raise "Backend not available".') + for device, reason in sorted(errors.items()): + print(f" {device:8} unavailable ({reason})") + if has_backend_conflict(): + print(" WARNING: OMP-offload is loaded alongside CPU/Thrust. Offload") + print(" regions will silently run on the host even though the") + print(" GPU query below succeeds. Give _core_gpu_omp_offload.so") + print(" a directory of its own and rebuild.") + + print("-" * 60) + print("Hardware detected (independent of what was compiled):") if info['gpu_available']: - print(f"✓ GPU Available: {info['gpu_type']}") - if info['gpu_count'] > 0: - print(f" GPU Count: {info['gpu_count']}") + print(f" GPU: {info['gpu_type']}" + + (f", count {info['gpu_count']}" if info['gpu_count'] > 0 else "")) else: - print("✗ No GPU detected") - - print(f"✓ CPU Available: Always") - print("="*60) + print(" GPU: none detected") + print(" CPU: always present") + print("=" * 60) diff --git a/python/examples/README.md b/python/examples/README.md index 4fce7c5..dcc7496 100644 --- a/python/examples/README.md +++ b/python/examples/README.md @@ -14,11 +14,14 @@ The standalone `run_sbd_diag.py` script needs nothing beyond what The SQD examples — `run_sqd_sbd.py` and `run_sqd_sbd.ipynb` — wrap SBD with the qiskit-addon-sqd self-consistent loop, which pulls in three -extra Python packages. Install them once into the same venv: +extra Python packages. Install them once into the same environment SBD +was built in: ```bash -source ~/venvs//bin/activate -pip install pyscf qiskit "qiskit-addon-sqd>=0.13.1" +conda activate sbd # the env from the Installation section of ../../README.md +pip install qiskit "qiskit-addon-sqd>=0.13.1" +# pyscf is already there if you used the conda recipe in ../../README.md; +# otherwise: conda install -y -c conda-forge pyscf ``` - **`pyscf`** — reads FCIDUMP, restores 4-fold integral symmetry. @@ -196,7 +199,9 @@ documented further here. ## Backend Selection -All compiled backends load eagerly at import. Select per-call via `--device`: +Every backend the toolchain supported was compiled into this one install, and +each is imported **lazily, on first use** — normally one per process. Select +per-call via `--device`: ```bash --device cpu # host OpenMP (default) @@ -205,18 +210,28 @@ All compiled backends load eagerly at import. Select per-call via `--device`: --device auto # GPU if available, else CPU ``` -`gpu-omp` links a different OpenMP runtime (`libnvomp`) than `cpu`/`gpu`, so it is -normally built into its own install — see the [Python Bindings README](../../README.md). -`sbd.available_backends()` reports what the current install actually has. +`sbd.available_backends()` reports what this install actually has (a static scan +— it does not import anything, so it is safe to call outside `mpirun`), and +`sbd.loaded_backends()` reports what the current process has pulled in. -Within Python, backends can also be switched at runtime without re-initialization: +Lazy loading is what makes the three backends safe to co-install: see +[Backend Architecture](../../README.md#backend-architecture) in the Python +Bindings README for why. One consequence is worth knowing when you write your +own driver — **do not import the CPU and `gpu-omp` backends into the same +process.** They share NVHPC's `libnvomp`, and loading `_core_cpu` first leaves +it initialised host-only, after which offload regions run on the host while +device queries still report a GPU. `sbd.has_backend_conflict()` returns True if +that has happened. Loading `cpu` and `gpu` (Thrust) together is fine. + +Within Python, backends can be selected per call — no re-initialization needed: ```python import sbd # No init() needed — auto-initializes on first call result_cpu = sbd.tpb_diag(..., device='cpu') -result_gpu = sbd.tpb_diag(..., device='gpu') +result_gpu = sbd.tpb_diag(..., device='gpu') # fine alongside 'cpu' +# result_omp = sbd.tpb_diag(..., device='gpu-omp') # NOT in the same process as 'cpu' ``` ## Available Test Data diff --git a/python/examples/run_sqd_sbd.ipynb b/python/examples/run_sqd_sbd.ipynb index d9e14a0..09bbd88 100644 --- a/python/examples/run_sqd_sbd.ipynb +++ b/python/examples/run_sqd_sbd.ipynb @@ -606,7 +606,7 @@ "- **Why this works in a notebook**: SBD is MPI-native at the C++ level, but a Jupyter kernel is a single Python process. `mpi4py` auto-initializes MPI with `MPI.COMM_WORLD` of size 1; SBD's collectives all become no-ops; the full SQD loop runs on rank 0.\n", "- **Why a curated counts file**: qiskit-addon-sqd postselects on (5α, 5β), which uniform-random strings almost never satisfy. `count_dict_h2o.json` supplies determinants that do. From iteration 2 on, configuration recovery refines them from the previous diagonalization's occupancies — no `initial_occupancies` needed.\n", "- **Multi-rank**: launch [`run_sqd_sbd.py`](./run_sqd_sbd.py) with `mpirun -np N python …` for production scale.\n", - "- **GPU**: pass `DeviceConfig.gpu()` (Thrust) or `DeviceConfig.gpu_omp()` (LLVM offload) on a CUDA-capable machine. Not available on macOS — keep `DeviceConfig.cpu()` here." + "- **GPU**: pass `DeviceConfig.gpu()` (Thrust) or `DeviceConfig.gpu_omp()` (NVHPC OpenMP target offload) on a CUDA-capable machine. Not available on macOS — keep `DeviceConfig.cpu()` here." ] } ], diff --git a/python/examples/run_sqd_sbd.py b/python/examples/run_sqd_sbd.py index 98ab41d..dc9eacb 100644 --- a/python/examples/run_sqd_sbd.py +++ b/python/examples/run_sqd_sbd.py @@ -68,7 +68,7 @@ def parse_args(): choices=["auto", "cpu", "gpu", "gpu-omp", "gpu-nvidia-omp"], default="cpu", help="cpu | gpu (NVHPC Thrust) | gpu-omp = gpu-nvidia-omp " - "(LLVM OpenMP-offload) | auto") + "(NVHPC OpenMP target offload) | auto") # SQD outer-loop parameters p.add_argument("--samples_per_batch", type=int, default=300) diff --git a/setup.py b/setup.py index ad09404..f0feade 100644 --- a/setup.py +++ b/setup.py @@ -14,6 +14,8 @@ import sys import os import subprocess +import re +import sysconfig import pybind11 @@ -35,24 +37,58 @@ def get_mpi4py_include(): return None -def _mpi_config_from_mpicc(): - """Probe the mpicc compiler wrapper for include/library/link flags. +def _mpi_prefix_from_mpi4py(): + """Derive the MPI install prefix from the library mpi4py is linked against. + + mpi4py is a build requirement, so it is importable here. Deriving the prefix + from it -- rather than from whichever mpicc happens to be first on PATH -- + guarantees the extensions link the same MPI that mpi4py uses. A mismatch is + not a build error: it surfaces at run time as an MPI_Init abort (e.g. MPICH + reporting "unsupported PMI version PMIx" under an Open MPI launcher), which + is considerably harder to diagnose. - Returns (include_dirs, library_dirs, libraries) or None if mpicc is - absent or does not understand the --showme flags (an OpenMPI-ism; - MPICH's wrapper does not support them). + Returns the prefix, or None when it cannot be determined -- a manylinux + wheel bundling its own MPI, a static build, or a platform with neither ldd + nor otool. Callers fall back to MPI_HOME. """ try: - compile_flags = subprocess.check_output(['mpicc', '--showme:compile'], - universal_newlines=True).strip().split() - link_flags = subprocess.check_output(['mpicc', '--showme:link'], - universal_newlines=True).strip().split() + import mpi4py + except ImportError: + return None + pkg_dir = os.path.dirname(mpi4py.__file__) + try: + exts = [n for n in sorted(os.listdir(pkg_dir)) + if n.startswith('MPI.') and n.endswith(('.so', '.dylib', '.pyd'))] + except OSError: + return None + # Some builds ship one extension per MPI flavour (MPI.mpich.*, MPI.openmpi.*) + # and choose at import time, so linkage cannot tell us which one is in use. + # Ambiguous: let the caller fall back to MPI_HOME. + if len(exts) != 1: + return None + ext = exts[0] + probe = ['otool', '-L'] if sys.platform == 'darwin' else ['ldd'] + try: + out = subprocess.check_output(probe + [os.path.join(pkg_dir, ext)], + universal_newlines=True, + stderr=subprocess.DEVNULL) except Exception: return None - include_dirs = [flag[2:] for flag in compile_flags if flag.startswith('-I')] - library_dirs = [flag[2:] for flag in link_flags if flag.startswith('-L')] - libraries = [flag[2:] for flag in link_flags if flag.startswith('-l')] - return include_dirs, library_dirs, libraries + # Take the resolved path and normalise it. A conda-installed MPI is reached + # through a relative path with '..' segments (mpi4py/../../../libmpi.so.12), + # so matching a literal '/lib/libmpi' misses it entirely; realpath also + # follows the libmpi.so.12 -> libmpi.so.12.x.y symlink chain. + match = re.search(r'=>\s*(\S*libmpi\S*)', out) or re.search(r'(\S*libmpi\S*)', out) + if not match: + return None + raw = match.group(1) + # macOS records @rpath/@loader_path-relative install names; resolving those + # means walking LC_RPATH, which is not worth it here -- MPI_HOME covers it. + if not raw.startswith('/'): + return None + lib_path = os.path.realpath(raw) + prefix = os.path.dirname(os.path.dirname(lib_path)) + return prefix if os.path.exists(os.path.join(prefix, 'include', 'mpi.h')) else None def _building_extensions(): @@ -71,55 +107,223 @@ def _building_extensions(): return not {'sdist', 'egg_info'}.intersection(sys.argv[1:]) +def _mpi_prefix_from_env_prefix(): + """MPI installed inside the active Python environment. + + Covers conda (`conda install mpich`/`openmpi`) and the PyPI `mpich`/`openmpi` + wheels, both of which drop mpi.h and libmpi straight into the environment + prefix. Needed because mpi4py's linkage is not always readable: a conda-forge + mpi4py ships one extension per MPI flavour (MPI.mpich.*, MPI.openmpi.*) so + linkage cannot say which is active, and on macOS the install name is + @rpath-relative. In both cases the prefix answers the question directly. + + This is a heuristic, not authoritative like mpi4py's own linkage, so callers + let MPI_HOME override it silently rather than treating a difference as a + conflict. + """ + for prefix in (sys.prefix, getattr(sys, 'base_prefix', sys.prefix)): + if prefix and os.path.exists(os.path.join(prefix, 'include', 'mpi.h')): + lib = os.path.join(prefix, 'lib') + if os.path.isdir(lib) and any(n.startswith('libmpi') for n in os.listdir(lib)): + return prefix + return None + + +def _mpi_config_from_mpicc(): + """Probe the mpicc wrapper for include/library/link flags. + + Needed because a prefix without $prefix/include/mpi.h is not necessarily the + wrong MPI: distros that split MPI into a -devel package put the headers + elsewhere (Debian's Open MPI uses /usr/lib//openmpi/include), and + mpicc is the thing that knows where its own headers live. + + `-show` prints the whole command line and is understood by both Open MPI and + MPICH -- checked against openmpi 5.0.10 and mpich 5.0.1. Open MPI's + `--showme:compile` is deliberately not tried: it yields nothing `-show` does + not, and MPICH's wrapper treats it as a source file and tries to compile it. + + Returns (include_dirs, library_dirs, libraries), or None if mpicc is absent or + prints no usable flags. + """ + try: + tokens = subprocess.check_output( + ['mpicc', '-show'], universal_newlines=True, + stderr=subprocess.DEVNULL).split() + except Exception: + return None + + inc, lib, libs = [], [], [] + for token in tokens: + for prefix, acc in (('-I', inc), ('-L', lib), ('-l', libs)): + if token.startswith(prefix) and token[2:] and token[2:] not in acc: + acc.append(token[2:]) # MPICH repeats -I/-L; collapse them + break + if not (inc or lib): + return None + # Open MPI names -lmpi, MPICH -lmpi -lpmpi; take what the wrapper says. + return inc, lib, libs or ['mpi'] + + def get_mpi_config(): - # Prefer MPI_HOME, but only trust it if mpi.h actually lives at - # $MPI_HOME/include. Distros that split MPI into a -devel package - # (e.g. Fedora's environment-modules sets MPI_HOME=/usr/lib64/openmpi - # while headers live in /usr/include/openmpi-x86_64) break the naive - # $MPI_HOME/include assumption, so we fall through to mpicc there. - mpi_home = os.environ.get('MPI_HOME', None) - if mpi_home: - mpi_include = os.path.join(mpi_home, 'include') - mpi_lib = os.path.join(mpi_home, 'lib') - if os.path.exists(os.path.join(mpi_include, 'mpi.h')): - print(f"Using MPI from MPI_HOME: {mpi_home}") - return [mpi_include], [mpi_lib], ['mpi'] - print(f"Notice: MPI_HOME={mpi_home} set but {mpi_include}/mpi.h " - "not found; falling back to mpicc detection.") + """Resolve MPI include/lib dirs, preferring the MPI that mpi4py uses. + + Order of precedence: MPI_HOME when set (explicit override, and the escape + hatch for layouts this cannot infer), otherwise the MPI mpi4py is linked + against. When both are known and disagree, that is a hard error -- linking a + different MPI than mpi4py aborts at MPI_Init, so failing the build is much + the cheaper outcome. + + Both MPICH and Open MPI install libmpi, so -lmpi is correct for either. + """ + derived = _mpi_prefix_from_mpi4py() + # Only mpi4py's own linkage is authoritative enough to contradict MPI_HOME. + from_prefix = None if derived else _mpi_prefix_from_env_prefix() + mpi_home = os.environ.get('MPI_HOME') or None + + if mpi_home and derived and \ + os.path.realpath(mpi_home) != os.path.realpath(derived): + print(f"Error: MPI_HOME={mpi_home} is not the MPI that mpi4py is linked " + f"against ({derived}).\n" + " Building against a different MPI than mpi4py aborts at " + "MPI_Init rather than at build time.\n" + " Either unset MPI_HOME to use mpi4py's MPI, or rebuild " + "mpi4py against MPI_HOME:\n" + f" MPICC={mpi_home}/bin/mpicc pip install --no-binary=mpi4py " + "--force-reinstall --no-cache-dir mpi4py\n" + " Then verify: python -c \"from mpi4py import MPI; " + "print(MPI.Get_library_version())\"") + sys.exit(1) + + prefix = mpi_home or derived or from_prefix + if prefix: + include_dir = os.path.join(prefix, 'include') + lib_dir = next((os.path.join(prefix, d) for d in ('lib', 'lib64') + if os.path.isdir(os.path.join(prefix, d))), + os.path.join(prefix, 'lib')) + source = ('MPI_HOME' if mpi_home else + 'mpi4py' if derived else 'the environment prefix') + if os.path.exists(os.path.join(include_dir, 'mpi.h')): + print(f"Using MPI from {source}: {prefix}") + return [include_dir], [lib_dir], ['mpi'] + # Same MPI, unusual layout -- ask mpicc before giving up on it. + print(f"Notice: no mpi.h under {prefix} (from {source}); asking mpicc.") + mpicc_config = _mpi_config_from_mpicc() + if mpicc_config is not None: + print("Using MPI detected from mpicc") + return mpicc_config + print(f"Warning: {include_dir}/mpi.h not found and mpicc unusable; " + "trying the prefix anyway. Check MPI_HOME points at an MPI " + "*prefix*, not its lib or bin directory.") + return [include_dir], [lib_dir], ['mpi'] mpicc_config = _mpi_config_from_mpicc() if mpicc_config is not None: print("Using MPI detected from mpicc") return mpicc_config - # Last resort: honor MPI_HOME even without a discoverable mpi.h, so a - # deliberately-set MPI_HOME on an unusual layout still gets a chance. - if mpi_home: - print(f"Warning: using MPI_HOME={mpi_home} despite missing " - f"{mpi_include}/mpi.h and unusable mpicc.") - return [os.path.join(mpi_home, 'include')], \ - [os.path.join(mpi_home, 'lib')], ['mpi'] - if not _building_extensions(): print("Notice: Could not detect MPI, but no extension is being " "compiled; continuing without MPI flags.") return [], [], ['mpi'] - print("Error: Could not detect MPI. Please set MPI_HOME environment " - "variable, or ensure mpicc is on PATH.") + print("Error: Could not determine which MPI to build against.\n" + " mpi4py did not reveal an MPI prefix containing include/mpi.h " + "(a wheel that bundles its own MPI will do that), so set MPI_HOME:\n" + " MPI_HOME=/path/to/mpi pip install -e . --no-build-isolation\n" + " Better, install mpi4py against that MPI first so the two " + "cannot diverge:\n" + " MPICC=/path/to/mpi/bin/mpicc pip install --no-binary=mpi4py " + "mpi4py") sys.exit(1) -def _resolve_gpu_arch(default='cc90'): - """Pick the nvc++ -gpu= value from env, with back-compat alias. +def _amdgpu_arch(compiler): + """Ask ROCm which GPU this machine has, e.g. 'gfx90a'. None if it cannot. + + Uses the amdgpu-arch that ships BESIDE the chosen compiler, never one found + on PATH. ROCm installs versioned trees side by side and also exposes + /usr/bin/amdgpu-arch via alternatives, so a PATH lookup can easily report + the arch from a different ROCm than the one doing the compiling. Note that + amdgpu-arch lives only in lib/llvm/bin, not in the prefix's bin, so resolve + the compiler symlink before looking next to it. + + Prints one line per GPU, so dedupe. Returns None on a GPU-less build host + (a container stage, a login node), which the caller turns into a request to + set SBD_GPU_ARCH explicitly. + """ + here = os.path.dirname(os.path.realpath(compiler)) + probe = os.path.join(here, 'amdgpu-arch') + if not os.path.exists(probe): + return None + try: + out = subprocess.check_output([probe], universal_newlines=True, + stderr=subprocess.DEVNULL, timeout=30) + except Exception: + return None + arches = sorted({ln.strip() for ln in out.splitlines() if ln.strip()}) + return ','.join(arches) or None + + +def _gpu_arch_flags(vendor, arch): + """Compiler flags that pin the GPU architecture, for compile AND link. + + Returned as a list so an unset arch contributes nothing at all. + + The two vendors spell this differently, and AMD cannot take a comma list: + nvc++ accepts -gpu=cc80,cc90,cc100 as one flag, while clang wants a repeated + --offload-arch=. Both MUST be passed at link as well as compile -- see the + long note on the Thrust extension below for what silently goes missing + otherwise. + """ + if not arch: + return [] + if vendor == 'amd': + return [f'--offload-arch={a}' for a in arch.split(',') if a] + return [f'-gpu={arch}'] + + +def _resolve_gpu_arch(vendor='nvidia', compiler=None): + """Return the GPU architecture to target, or None to let the toolchain pick. + + OPTIONAL BY DESIGN. + + On AMD the arch is auto-DETECTED rather than left implicit: unset, we ask + ROCm's amdgpu-arch what this machine has (e.g. gfx90a for MI250X, gfx942 for + MI300X) and pin that. amdclang++ has no useful built-in default the way + nvc++ does, so if detection fails -- a build host with no GPU -- SBD_GPU_ARCH + becomes mandatory and the build stops with that message. Several + architectures can be named at once (gfx90a,gfx942); each becomes its own + --offload-arch flag. As on NVIDIA there is no JIT fallback, so an + architecture that is not listed will not run. + + Verify what actually landed rather than trusting the flag: + + llvm-objdump --offloading + + The NVIDIA contract is unchanged: + + * Set (e.g. ``cc90``, or ``cc80,cc90,cc100`` for a portable + multi-architecture binary): honored exactly, and passed at BOTH compile + and link. Passing it at link is not optional -- the device-link step is + where the final SASS is generated, so a value given only at compile time + is discarded and you silently get the toolchain's own target instead. + * Unset: no ``-gpu=`` flag is emitted at all and nvc++ targets the GPU of + the machine the toolchain was installed on. That is the right default for + a local build on the machine you will run on, and it keeps a plain + ``pip install`` working with no environment setup. + + Build a multi-architecture binary whenever the artifact travels -- a + container image, a shared filesystem, a wheel for a mixed cluster. These + builds embed no PTX, so an architecture that is not listed has no JIT + fallback: it simply will not run. ``ccall-major`` covers one target per + major generation if you would rather not enumerate. - Both Thrust (_core_gpu_thrust) and OMP-offload (_core_gpu_omp_offload) - backends compile with nvc++ and take the same -gpu= flag, so we - use a single env var. + Verify what actually landed rather than trusting the flag: - Reads SBD_GPU_ARCH (canonical name since v1.6). Falls back to the - deprecated SBD_GPU_ARCH_NVIDIA (v1.5 and earlier) with a notice so - existing setup scripts keep working through the transition. + cuobjdump --list-elf + + Reads SBD_GPU_ARCH, honoring the deprecated SBD_GPU_ARCH_NVIDIA (v1.5 and + earlier) with a notice so existing scripts keep working. """ val = os.environ.get('SBD_GPU_ARCH') if val: @@ -128,35 +332,67 @@ def _resolve_gpu_arch(default='cc90'): if legacy: print(f"Notice: SBD_GPU_ARCH_NVIDIA={legacy!r} is deprecated since " "v1.6 (single SBD_GPU_ARCH covers both Thrust and OMP-offload " - "since LLVM/clang was removed). Honoring it as a back-compat " + "now that the LLVM path is gone). Honoring it as a back-compat " "alias. Please switch to SBD_GPU_ARCH.") return legacy - return default + + if vendor == 'amd': + detected = _amdgpu_arch(compiler) if compiler else None + if detected: + print(f"Notice: SBD_GPU_ARCH is not set; detected {detected} via " + "amdgpu-arch and\n" + " targeting exactly that. If this artifact will run " + "anywhere else --\n" + " a container image, a shared filesystem, a mixed-GPU " + "cluster -- set\n" + " SBD_GPU_ARCH to every architecture you need, e.g. " + "gfx90a,gfx942.\n" + " There is no JIT fallback, so an unlisted " + "architecture cannot run.") + return detected + # The caller turns this into a hard error. Not a silent pass: unlike + # nvc++, amdclang++ has no built-in default worth inheriting, so a build + # with no arch at all produces a module that runs nowhere. + print("Notice: SBD_GPU_ARCH is not set and amdgpu-arch could not " + "report an\n" + " architecture (normal on a build host with no AMD GPU).") + return None + + print("Notice: SBD_GPU_ARCH is not set; letting nvc++ target the GPU of the\n" + " machine this toolchain was installed on. Fine for a local\n" + " build. If this artifact will run anywhere else -- a container\n" + " image, a shared filesystem, a mixed-GPU cluster -- set it to\n" + " every architecture you need, e.g.\n" + " SBD_GPU_ARCH=cc80,cc90,cc100 (A100 / H100 / GB200-B200)\n" + " SBD_GPU_ARCH=ccall-major (one per major generation)\n" + " No PTX is embedded, so an unlisted architecture cannot JIT.") + return None -def _route_build_through_nvhpc(nvc_path): - """Configure distutils + sysconfig so a setup() call uses nvc++. +def _route_build_through_gpu_compiler(gpu_path, vendor='nvidia'): + """Configure distutils + sysconfig so a setup() call uses the GPU compiler. - Called by both the Thrust and OMP-offload extension blocks (both - compile with nvc++). Idempotent — second call is a no-op. + Called by the extension blocks that need it: Thrust and OMP-offload on + NVIDIA (both nvc++), OMP-offload on AMD (amdclang++). Idempotent — second + call is a no-op. Effect: distutils' UnixCCompiler will pick up CC/CXX/LDSHARED from os.environ and use them for every Extension in this setup() call. Also clears CFLAGS/CXXFLAGS/CPPFLAGS and rewrites sysconfig to drop - gcc-specific tokens nvc++ rejects (RHEL 9 CPython injects a long - list — see comment below). + tokens the chosen compiler rejects (see the two lists below). - Co-builds with the CPU extension are safe: nvc++ accepts the CPU - block's `-fopenmp -O3 -std=c++17` flags (treats -fopenmp as -mp). + Co-builds with the CPU extension are safe under either vendor: both nvc++ + and amdclang++ accept the CPU block's `-fopenmp -O3 -std=c++17` (nvc++ + treats -fopenmp as -mp; amdclang++ IS clang, so it takes them natively). """ - if os.environ.get('_SBD_NVHPC_ROUTING_APPLIED'): + if os.environ.get('_SBD_GPU_ROUTING_APPLIED'): return - os.environ['_SBD_NVHPC_ROUTING_APPLIED'] = '1' + os.environ['_SBD_GPU_ROUTING_APPLIED'] = '1' - # Respect user-set CC/CXX (e.g. cross-toolchain); otherwise pin nvc++. - os.environ.setdefault('CC', nvc_path) - os.environ.setdefault('CXX', nvc_path) - os.environ.setdefault('LDSHARED', f'{nvc_path} -shared') + # Respect user-set CC/CXX (e.g. cross-toolchain); otherwise pin the GPU one. + os.environ.setdefault('CC', gpu_path) + os.environ.setdefault('CXX', gpu_path) + os.environ.setdefault('LDSHARED', f'{gpu_path} -shared') os.environ.setdefault('CFLAGS', '') os.environ.setdefault('CXXFLAGS', '') os.environ.setdefault('CPPFLAGS', '') @@ -169,27 +405,54 @@ def _route_build_through_nvhpc(nvc_path): # (requires v3+). distutils pulls these from sysconfig in addition # to os.environ.CFLAGS, so blanking the latter alone is not enough # — we rewrite the sysconfig dict itself. - import sysconfig, re as _re + # + # amdclang++ needs FAR less scrubbing, because it is clang and accepts the + # gcc spellings. Measured against ROCm 10.0 / AMD clang 23 with + # --offload-arch=gfx90a, every token above compiles clean EXCEPT + # -fcf-protection, which is rejected as "option 'cf-protection=return' + # cannot be specified on this target" -- the flag is applied to the amdgcn + # device pass too, and there it is meaningless. -march=x86-64-v2 is fine for + # clang and is deliberately NOT rewritten to v3 here; that rewrite exists + # only because nvc++ requires v3+. _cfg = sysconfig.get_config_vars() - _strip_tokens = ( - '-grecord-gcc-switches', - '-Wp,-D_FORTIFY_SOURCE=2', - '-Wp,-D_GLIBCXX_ASSERTIONS', - '-fstack-protector-strong', - '-fasynchronous-unwind-tables', - '-fstack-clash-protection', - '-fcf-protection', - '-fwrapv', - '-Wno-unused-result', - ) + if vendor == 'amd': + _strip_tokens = ( + '-fcf-protection', + ) + else: + _strip_tokens = ( + '-grecord-gcc-switches', + '-Wp,-D_FORTIFY_SOURCE=2', + '-Wp,-D_GLIBCXX_ASSERTIONS', + '-fstack-protector-strong', + '-fasynchronous-unwind-tables', + '-fstack-clash-protection', + '-fcf-protection', + '-fwrapv', + '-Wno-unused-result', + ) for _k in list(_cfg.keys()): _v = _cfg[_k] if not isinstance(_v, str): continue for _bad in _strip_tokens: _v = _v.replace(_bad, '') - _v = _v.replace('-march=x86-64-v2', '-march=x86-64-v3') - _cfg[_k] = _re.sub(r' +', ' ', _v).strip() + # nvc++ only: it rejects x86-64-v2 and requires v3+. clang accepts v2, + # so leave it alone there rather than silently raising the CPU baseline + # of the AMD build above what the caller's Python asked for. + if vendor != 'amd': + _v = _v.replace('-march=x86-64-v2', '-march=x86-64-v3') + # conda's Python bakes '-B $CONDA_PREFIX/compiler_compat' into + # CC/CXX/LDSHARED/LDCXXSHARED. nvc++ rejects -B and hands the + # path to the linker as an input file, so drop just that flag + # and keep conda's -L/-rpath entries intact. + # + # Dropped for amdclang++ too, for a different reason: clang accepts -B + # perfectly well, but that directory holds conda's own (old) `ld`, and + # the offload link runs through clang-linker-wrapper -> ld.lld. Letting + # -B redirect the linker there invites a mismatch for no benefit. + _v = re.sub(r'-B\s*\S*compiler_compat\S*', '', _v) + _cfg[_k] = re.sub(r' +', ' ', _v).strip() def find_nvidia_hpc_sdk(): @@ -213,6 +476,92 @@ def find_nvidia_hpc_sdk(): return None, False +def find_rocm_toolchain(): + """Locate amdclang++ for the AMD OpenMP target-offload backend. + + Deliberately the same shape as find_nvidia_hpc_sdk(): an explicit env var, + else PATH. ROCM_HOME is this project's knob, matching NVHPC_HOME, and its + job is picking a specific ROCm on a node with several installed -- common, + since ROCm ships side-by-side versioned trees. + + Neither variable is required. A PATH lookup already covers both the + module-based case (`module add rocm/` prepends its bin) and a stock + install (ROCm's packages leave amdclang++ in /usr/bin via alternatives), and + not every cluster provides modules. + + amdclang++ IS LLVM clang -- ROCm ships it with the amdgcn OpenMP offload + runtime and matching device libraries already built, so nothing has to be + compiled from source to get offload working. Unlike the NVHPC branch there + is no need to touch PATH: amdclang++ finds its device libraries relative to + its own InstalledDir, so invoking it by absolute path is enough. + """ + import shutil + # ROCM_HOME (this project's knob, matching NVHPC_HOME) wins over ROCM_PATH + # (ROCm's own variable, usually set by a module file). + for var in ('ROCM_HOME', 'ROCM_PATH'): + rocm_home = os.environ.get(var) or None + if not rocm_home: + continue + # $ROCM_HOME/bin/amdclang++ is normally a symlink to the second path; + # older layouts only have lib/llvm/bin. + for rel in ('bin/amdclang++', 'lib/llvm/bin/amdclang++'): + cand = os.path.join(rocm_home, rel) + if os.path.exists(cand): + print(f"Found ROCm at: {rocm_home} (via {var})") + return cand, True + print(f"Warning: {var} set to {rocm_home} but amdclang++ not found") + amdcxx_path = shutil.which('amdclang++') + if amdcxx_path: + print(f"Found amdclang++ in PATH: {amdcxx_path}") + return amdcxx_path, True + return None, False + + +def detect_gpu_toolchain(): + """Pick the GPU toolchain to build with: ('nvidia'|'amd'|None, compiler). + + NVHPC is probed first only because it is the long-established path here; a + machine with exactly one GPU toolchain installed gets that one either way. + SBD_GPU_VENDOR forces the choice for the rare host carrying both (a build + node serving a mixed cluster), where the auto-answer would otherwise be an + accident of probe order. + + Note the asymmetry in what each vendor can build: nvc++ drives BOTH the + Thrust and the OpenMP-offload backends, whereas ROCm drives OpenMP offload + only. Upstream SBD's Thrust path is wired to nvc++ flags (-cuda, -gpu=), so + there is no rocThrust configuration to build even though some HIP scaffolding + exists upstream. + """ + forced = (os.environ.get('SBD_GPU_VENDOR') or '').strip().lower() + if forced not in ('', 'nvidia', 'amd', 'none'): + print(f"Error: Invalid SBD_GPU_VENDOR={forced!r}. " + "Valid values: nvidia, amd, none") + sys.exit(1) + if forced == 'none': + print("SBD_GPU_VENDOR=none - skipping GPU toolchain detection") + return None, None + + if forced != 'amd': + nvcxx, ok = find_nvidia_hpc_sdk() + if ok: + return 'nvidia', nvcxx + if forced == 'nvidia': + print("Error: SBD_GPU_VENDOR=nvidia but nvc++ was not found. " + "Set NVHPC_HOME.") + sys.exit(1) + + if forced != 'nvidia': + amdcxx, ok = find_rocm_toolchain() + if ok: + return 'amd', amdcxx + if forced == 'amd': + print("Error: SBD_GPU_VENDOR=amd but amdclang++ was not found. " + "Set ROCM_HOME, or put ROCm's bin directory on PATH.") + sys.exit(1) + + return None, None + + # Get MPI configuration mpi_includes, mpi_lib_dirs, mpi_libs = get_mpi_config() @@ -249,67 +598,129 @@ def find_nvidia_hpc_sdk(): # RPATH so libraries are found at runtime without LD_LIBRARY_PATH extra_link_args = ['-fopenmp'] +_rpath_dirs = [] for lib_dir in library_dirs: - extra_link_args.append(f'-Wl,--rpath,{lib_dir}') + if lib_dir not in _rpath_dirs: # a dir may appear as both MPI and BLAS + _rpath_dirs.append(lib_dir) + extra_link_args.append(f'-Wl,--rpath,{lib_dir}') print(f"RPATH will be set to: {library_dirs}") -# Detect NVHPC. nvc++ is shared between two GPU backends here: -# 1. _core_gpu_thrust (Thrust + CUDA path, nvc++ -cuda) -# 2. _core_gpu_omp_offload (OpenMP target offload, nvc++ -mp=gpu) -gpu_compiler, has_nvhpc = find_nvidia_hpc_sdk() +# Runtime search order matters: conda's Python injects -Wl,-rpath,$CONDA_PREFIX/lib +# into LDSHARED/LDCXXSHARED, and setuptools places those flags *before* the ones +# built above. A conda-installed library therefore wins over an explicitly +# requested one -- BLAS_LIB_PATH gets honored at link time and silently ignored at +# run time, which is how you end up running conda's generic OpenBLAS while +# believing you selected a tuned build. +# +# Demote conda: drop its rpath entries (keeping its -L, so link-time discovery of +# conda-provided libraries still works) and re-add the directory last, as a +# fallback behind anything the caller asked for. +_conda_prefix = os.environ.get('CONDA_PREFIX') +if _conda_prefix: + _conda_lib = os.path.join(_conda_prefix, 'lib') + _scfg = sysconfig.get_config_vars() + _rpath_re = re.compile(r'-Wl,-rpath(?:-link)?,' + re.escape(_conda_lib) + r'(?=\s|$)') + for _key in ('LDSHARED', 'LDCXXSHARED'): + _val = _scfg.get(_key) + if isinstance(_val, str): + _scfg[_key] = re.sub(r' +', ' ', _rpath_re.sub('', _val)).strip() + if _conda_lib not in _rpath_dirs: # skip if already requested explicitly + _rpath_dirs.append(_conda_lib) + extra_link_args.append(f'-Wl,--rpath,{_conda_lib}') + print(f"RPATH fallback appended last: {_conda_lib}") + +# Detect the GPU toolchain. What it can build depends on the vendor: +# NVIDIA (nvc++) 1. _core_gpu_thrust (Thrust + CUDA, nvc++ -cuda) +# 2. _core_gpu_omp_offload (OMP offload, nvc++ -mp=gpu) +# AMD (amdclang++) _core_gpu_omp_offload (OMP offload, --offload-arch) +# +# The OMP-offload backend is ONE module and ONE device string ('gpu-omp') for +# both vendors: it is the same bindings.cpp with the same USE_GPU + +# USE_OMP_OFFLOAD macros, just a different compiler driving it. A given install +# serves one GPU vendor -- no wheels are published, every install compiles on the +# target machine -- so a vendor-suffixed second device string would buy nothing +# and would undo the deprecation of gpu_nvidia_omp() in favour of gpu_omp(). +# Which vendor a build targeted is recorded on the module as +# __sbd_offload_target__ (e.g. 'amdgcn-amd-amdhsa:gfx90a') so it stays +# introspectable. +gpu_vendor, gpu_compiler = detect_gpu_toolchain() +has_gpu_toolchain = gpu_compiler is not None +# Only NVHPC can build the Thrust backend: upstream SBD wires that path to nvc++ +# flags (-cuda, -gpu=), with no rocThrust configuration. +has_nvhpc = gpu_vendor == 'nvidia' # Determine which backends to build. -# auto : cpu + thrust GPU (if nvc++ present) +# auto : cpu, plus both GPU backends when nvc++ is present +# all : same as auto, but an error when nvc++ is missing # cpu : cpu only # gpu | gpu_thrust : thrust GPU only -# both : cpu + thrust # gpu_omp_offload : OpenMP target offload only (nvc++ -mp=gpu) +# `both` (cpu + thrust) was removed: it is a strict subset of `auto`, and the +# reason to keep the GPU backends apart went away with lazy loading. # -# gpu_omp_offload is built ALONE — it uses a different OpenMP runtime -# (libnvomp) than cpu (libgomp/libomp) and Thrust GPU (CPU OMP via -mp), -# and loading two backends with different OMP runtimes in one Python -# process produces "Another OpenMP runtime library has been detected" -# warnings and can deadlock at first OMP region. Build it into its own -# venv / install dir. +# All three may now be installed side by side. They used to be kept apart on +# the theory that they link different OpenMP runtimes; that is not what ldd +# shows -- when NVHPC is present every extension is compiled by nvc++ and all +# three link libnvomp. The real hazard was that _core_cpu, built without +# -mp=gpu, leaves that shared runtime initialised host-only, after which the +# OMP-offload backend cannot acquire a device and silently runs its target +# regions on the host (correct energies, exit 0, GPU still reported). Since the +# Python package now imports backends lazily -- one per process, on first use -- +# co-resident .so files no longer interfere, so `auto` builds everything the +# toolchain supports. build_backend = os.environ.get('SBD_BUILD_BACKEND', 'auto').lower() build_cpu = False build_gpu_thrust = False build_gpu_omp_offload = False -if build_backend == 'auto': +if build_backend in ('auto', 'all'): build_cpu = True build_gpu_thrust = has_nvhpc - if build_gpu_thrust: - print("\nAuto-detected nvc++ - will build both CPU and Thrust GPU backends") + build_gpu_omp_offload = has_gpu_toolchain + if has_nvhpc: + print("\nAuto-detected nvc++ - will build CPU, Thrust GPU and " + "OMP-offload GPU backends") + elif gpu_vendor == 'amd': + # No Thrust here, so `auto` yields two backends rather than three. + print("\nAuto-detected amdclang++ - will build CPU and OMP-offload GPU " + "backends (Thrust is NVIDIA-only)") else: - print("\nnvc++ not found - will build CPU backend only") + print("\nNo GPU compiler found - will build CPU backend only") + if build_backend == 'all' and not has_gpu_toolchain: + print("Error: SBD_BUILD_BACKEND=all requires a GPU toolchain " + "(NVHPC_HOME / nvc++, or ROCM_HOME / amdclang++).") + sys.exit(1) elif build_backend == 'cpu': build_cpu = True print("\nBuilding CPU backend only (SBD_BUILD_BACKEND=cpu)") elif build_backend in ('gpu', 'gpu_thrust'): build_gpu_thrust = True print(f"\nBuilding Thrust GPU backend only (SBD_BUILD_BACKEND={build_backend})") - if not has_nvhpc: - print("Warning: nvc++ not found, GPU build may fail") -elif build_backend == 'both': - build_cpu = True - build_gpu_thrust = True - print("\nBuilding both CPU and Thrust GPU backends (SBD_BUILD_BACKEND=both)") + if gpu_vendor == 'amd': + # Fail rather than warn: on AMD this is not a maybe-it-links situation, + # there is no rocThrust configuration to build at all. Silently falling + # back to CPU under a name that says 'gpu' is exactly the confusion the + # AMD path is meant to remove. + print("Error: the Thrust backend is NVIDIA-only (upstream wires it to " + "nvc++ -cuda).\n" + " On AMD use SBD_BUILD_BACKEND=gpu_omp_offload, or leave it " + "unset for CPU + OMP-offload.") + sys.exit(1) if not has_nvhpc: print("Warning: nvc++ not found, GPU build may fail") elif build_backend == 'gpu_omp_offload': - # Stand-alone build: this mode only emits _core_gpu_omp_offload.so. - # See note above on the OpenMP-runtime exclusivity constraint. build_gpu_omp_offload = True print("\nBuilding GPU OpenMP target-offload backend only " "(SBD_BUILD_BACKEND=gpu_omp_offload)") - if not has_nvhpc: - print("Error: gpu_omp_offload requires NVHPC_HOME / nvc++.") + if not has_gpu_toolchain: + print("Error: gpu_omp_offload requires a GPU toolchain: NVHPC_HOME / " + "nvc++, or ROCM_HOME / amdclang++.") sys.exit(1) else: print(f"Error: Invalid SBD_BUILD_BACKEND='{build_backend}'") - print("Valid values: auto, cpu, gpu (alias gpu_thrust), both, gpu_omp_offload") + print("Valid values: auto (= all backends the toolchain supports), all, " + "cpu, gpu (alias gpu_thrust), gpu_omp_offload") sys.exit(1) ext_modules = [] @@ -318,9 +729,49 @@ def find_nvidia_hpc_sdk(): print("\nConfiguring CPU backend (_core_cpu)") import platform if platform.system() == 'Darwin': - omp_inc = '/opt/homebrew/opt/libomp/include' - omp_lib = '/opt/homebrew/opt/libomp/lib' - openblas_lib = '/opt/homebrew/opt/openblas/lib' + # macOS has no system OpenMP, so libomp comes from a package manager. + # Prefer the conda env when it has one: those are the libraries actually + # LOADED at import time (resolved via the python executable's + # @loader_path/../lib), so building against Homebrew's copies instead + # means compiling against different libraries than the process runs on. + conda_prefix = os.environ.get('CONDA_PREFIX') + if conda_prefix and os.path.exists( + os.path.join(conda_prefix, 'include', 'omp.h')): + omp_inc = os.path.join(conda_prefix, 'include') + omp_lib = openblas_lib = os.path.join(conda_prefix, 'lib') + print(f"Darwin: libomp and BLAS from conda env {conda_prefix}") + else: + omp_inc = '/opt/homebrew/opt/libomp/include' + omp_lib = '/opt/homebrew/opt/libomp/lib' + openblas_lib = '/opt/homebrew/opt/openblas/lib' + if not os.path.exists(os.path.join(omp_inc, 'omp.h')): + # Fail here with the fix, rather than 100 lines later with + # "'omp.h' file not found" from the middle of a compile. + print("Error: no OpenMP runtime found on this macOS host.\n" + f" Looked in $CONDA_PREFIX/include and {omp_inc}.\n" + " Apple clang ships without OpenMP, so install one:\n" + " conda install -c conda-forge llvm-openmp (preferred)\n" + " brew install libomp") + sys.exit(1) + print("Darwin: libomp and BLAS from Homebrew (no conda libomp found)") + + # Say WHICH clang is compiling. distutils takes CC/CXX from the + # environment, else from sysconfig -- where conda records a bare + # 'clang++' that is resolved through PATH, so a Homebrew LLVM silently + # wins over both Apple clang and a conda toolchain. Printing it is the + # difference between a reproducible build and a mystery. + import shutil + _cxx = (os.environ.get('CXX') or sysconfig.get_config_var('CXX') + or 'clang++').split()[0] + _cxx_path = shutil.which(_cxx) or _cxx + try: + _cxx_ver = subprocess.check_output( + [_cxx_path, '--version'], universal_newlines=True, + stderr=subprocess.STDOUT).splitlines()[0] + except Exception: + _cxx_ver = '(version unknown)' + print(f"Darwin C++ compiler: {_cxx_path}\n" + f" {_cxx_ver} (pin it with CC/CXX)") cpu_compile_args = [ '-DSBD_TRADMODE', '-std=c++17', '-Xpreprocessor', '-fopenmp', '-O3', @@ -365,10 +816,17 @@ def find_nvidia_hpc_sdk(): print(f"Using compiler: {gpu_compiler}") # Auto-route the build through nvc++ + sanitize sysconfig flags. # No-op if the user already set CC/CXX manually. - _route_build_through_nvhpc(gpu_compiler) - gpu_arch = _resolve_gpu_arch(default='cc90') + _route_build_through_gpu_compiler(gpu_compiler, 'nvidia') + gpu_arch = _resolve_gpu_arch('nvidia', gpu_compiler) + # Emitted only when the user asked for a specific arch; otherwise omitted + # entirely so nvc++ picks the build machine's GPU (see _resolve_gpu_arch). + gpu_arch_flags = _gpu_arch_flags('nvidia', gpu_arch) print(f"NVHPC -gpu= arch: {gpu_arch} (set SBD_GPU_ARCH to override; " "nvc++ accepts cc and sm_)") + # Stamped for the same reason as the offload backend: so a built module can + # be asked what it targets. Unambiguously NVIDIA, but the architecture is + # not, and an arch mismatch is the usual reason a module refuses to run. + thrust_target = f'cuda:{gpu_arch or "toolchain-default"}' gpu_thrust_ext = Extension( 'sbd._core_gpu_thrust', @@ -387,8 +845,9 @@ def find_nvidia_hpc_sdk(): '--diag_suppress=declared_but_not_referenced,set_but_not_used', '-fmax-errors=0', '-fPIC', - f'-gpu={gpu_arch}', + *gpu_arch_flags, '-DSBD_MODULE_NAME=_core_gpu_thrust', + f'-DSBD_OFFLOAD_TARGET="{thrust_target}"', ], # NOTE: -cudalib (no value) makes nvc++ blanket-link every CUDA # library NVHPC ships, including math libs SBD never calls @@ -398,10 +857,27 @@ def find_nvidia_hpc_sdk(): # -lcublasmp" etc. SBD's GPU path only needs the CUDA runtime, so # explicitly link -lcudart instead. # - # -gpu= is repeated at link because the device-link step generates the - # final SASS: without it nvc++ silently targets its own default instead - # of the requested arch(es). - extra_link_args=extra_link_args + ['-mp', '-cuda', f'-gpu={gpu_arch}', + # -gpu= MUST be repeated here, at LINK time. The compile step above + # emits device code for every architecture in the list, but the + # device-link step decides what actually lands in the fatbin, and + # without -gpu= nvc++ keeps only its own built-in default and silently + # discards the rest -- no warning, exit 0. The result is a .so that + # runs only on whatever architecture that default happens to be. + # + # Measured on NVHPC 26.1 (aarch64) with SBD_GPU_ARCH=cc80,cc90,cc100: + # object after compile sm_80 sm_90 sm_100 + # .so linked without -gpu= sm_100 <- two arches lost + # .so linked with -gpu= sm_80 sm_90 sm_100 + # The default is compiled into nvc++, not detected from the hardware, + # so a GPU-less build host (a container stage, say) does not change it. + # + # The OMP-offload extension below already passes -gpu= at link, which is + # why only the Thrust backend was affected: a multi-arch build appeared + # to succeed and then failed on any GPU other than the build toolchain's + # default -- observed as a Thrust-only failure on H100 from a fatbin + # built for cc80,cc90,cc100. These builds embed no PTX, so there is no + # JIT fallback to mask it. + extra_link_args=extra_link_args + ['-mp', '-cuda', *gpu_arch_flags, '-lcudart'], ) ext_modules.append(gpu_thrust_ext) @@ -409,39 +885,92 @@ def find_nvidia_hpc_sdk(): if build_gpu_omp_offload: print("\nConfiguring GPU OpenMP target-offload backend (_core_gpu_omp_offload)") - print(f"Using compiler: {gpu_compiler}") - # Auto-route the build through nvc++ + sanitize sysconfig flags. - _route_build_through_nvhpc(gpu_compiler) - offload_arch = _resolve_gpu_arch(default='cc90') - print(f"NVHPC -gpu= arch: {offload_arch} (set SBD_GPU_ARCH to override)") - - gpu_omp_offload_ext = Extension( - 'sbd._core_gpu_omp_offload', - ['python/bindings.cpp'], - include_dirs=include_dirs, - libraries=libraries, - library_dirs=library_dirs, - language='c++', - extra_compile_args=[ + print(f"Using compiler: {gpu_compiler} (vendor: {gpu_vendor})") + # Auto-route the build through the GPU compiler + sanitize sysconfig flags. + _route_build_through_gpu_compiler(gpu_compiler, gpu_vendor) + offload_arch = _resolve_gpu_arch(gpu_vendor, gpu_compiler) + offload_arch_flags = _gpu_arch_flags(gpu_vendor, offload_arch) + + if gpu_vendor == 'amd': + print(f"ROCm offload arch: {offload_arch} " + "(set SBD_GPU_ARCH to override, e.g. gfx90a or gfx90a,gfx942)") + if not offload_arch: + # No baked-in default exists for amdclang++, and an unpinned build + # would produce a module that cannot run anywhere. + print("Error: could not detect the AMD GPU architecture and " + "SBD_GPU_ARCH is not set.\n" + " Set it explicitly, e.g. SBD_GPU_ARCH=gfx90a " + "(MI250X) or gfx942 (MI300X).\n" + " `amdgpu-arch` on a machine with the target GPU " + "prints the right value.") + sys.exit(1) + offload_target = f'amdgcn-amd-amdhsa:{offload_arch}' + # -fopenmp-offload-mandatory: refuse to emit a host fallback path at + # COMPILE time. It pairs with the runtime OMP_TARGET_OFFLOAD=MANDATORY + # that __init__.py sets before importing this backend; together they + # make a device-less rank a loud failure rather than a silent host run + # returning a plausible energy. + # + # No sbd_nvhpc_compat.h here: that shim exists because nvc++ lowers + # __builtin_ffsl inside `declare target` to a host-only symbol. clang + # lowers those builtins to device intrinsics natively, and the shim is + # #ifdef __NVCOMPILER anyway, so including it would be a no-op. + offload_compile_args = [ + '-O3', '-std=c++17', '-fPIC', + '-fopenmp', '-fopenmp-targets=amdgcn-amd-amdhsa', + *offload_arch_flags, + '-fopenmp-offload-mandatory', + '-DSBD_TRADMODE', + '-DUSE_GPU', + '-DUSE_OMP_OFFLOAD', + '-DOMPI_SKIP_MPICXX', + '-DSBD_MODULE_NAME=_core_gpu_omp_offload', + f'-DSBD_OFFLOAD_TARGET="{offload_target}"', + # Selects the AMD device-visibility variables in bindings.cpp, so + # the NVIDIA build keeps reading CUDA_VISIBLE_DEVICES first exactly + # as before. + '-DSBD_OFFLOAD_VENDOR_AMD', + # Upstream headers are template-heavy and noisy under clang; these + # are style warnings in vendored code, not actionable here. + '-Wno-sign-compare', '-Wno-unused-variable', + ] + offload_link_args = extra_link_args + [ + '-fopenmp', *offload_arch_flags, + ] + else: + print(f"NVHPC -gpu= arch: {offload_arch} (set SBD_GPU_ARCH to override)") + offload_target = f'nvptx64-nvidia-cuda:{offload_arch or "toolchain-default"}' + offload_compile_args = [ '-O3', '-std=c++17', '-fPIC', '-mp=gpu', - f'-gpu={offload_arch}', + *offload_arch_flags, '-Minfo=mp', '-DSBD_TRADMODE', '-DUSE_GPU', '-DUSE_OMP_OFFLOAD', '-DOMPI_SKIP_MPICXX', '-DSBD_MODULE_NAME=_core_gpu_omp_offload', + f'-DSBD_OFFLOAD_TARGET="{offload_target}"', # Force-include nvc++ shim so __builtin_ffsl / __builtin_popcountl # inside #pragma omp declare target lower to portable inlines # rather than __blt_pgi_ffsl (host-only NVHPC symbol that nvlink # can't resolve from device code). '-include', 'python/sbd_nvhpc_compat.h', - ], - extra_link_args=extra_link_args + [ + ] + offload_link_args = extra_link_args + [ '-mp=gpu', - f'-gpu={offload_arch}', - ], + *offload_arch_flags, + ] + + gpu_omp_offload_ext = Extension( + 'sbd._core_gpu_omp_offload', + ['python/bindings.cpp'], + include_dirs=include_dirs, + libraries=libraries, + library_dirs=library_dirs, + language='c++', + extra_compile_args=offload_compile_args, + extra_link_args=offload_link_args, ) ext_modules.append(gpu_omp_offload_ext) @@ -459,5 +988,6 @@ def find_nvidia_hpc_sdk(): if build_gpu_thrust: print(" - Thrust GPU backend: sbd._core_gpu_thrust") if build_gpu_omp_offload: - print(" - OpenMP-offload GPU backend: sbd._core_gpu_omp_offload") + print(" - OpenMP-offload GPU backend: sbd._core_gpu_omp_offload" + f" ({offload_target})") print() diff --git a/tox.ini b/tox.ini index 75068c6..dd3ecb9 100644 --- a/tox.ini +++ b/tox.ini @@ -17,6 +17,9 @@ passenv = BLAS_LIB_PATH BLAS_LIBS NVHPC_HOME + ROCM_HOME + ROCM_PATH + SBD_GPU_VENDOR SBD_GPU_ARCH SBD_BUILD_BACKEND SBD_TEST_DEVICE