feat(distributed): MPI/distributed bindings for Ginkgo

CMakeLists.txt

@@ -16,7 +16,23 @@ set(PYTHON_EXECUTABLE ${Python_EXECUTABLE})
 # copy all generated files to pyGinkgo that makes testing and installing easier
 set(CMAKE_LIBRARY_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/${PROJECT_NAME})
 
-# Find OpenMPI find_package(MPI REQUIRED)
+# Optional MPI / distributed support.
+#
+# When enabled:
+#   * find_package(MPI REQUIRED) is invoked
+#   * mpi4py headers are located via the active Python interpreter
+#   * GINKGO_BUILD_MPI=ON is propagated to the bundled Ginkgo (FetchContent),
+#     and a pre-installed Ginkgo without MPI support is rejected
+#   * the PYGINKGO_BUILD_MPI compile definition is set, gating all
+#     #ifdef-protected distributed source files
+#   * the MPI implementation flavor and library version are baked into a
+#     generated header so the import-time check can verify mpi4py matches
+option(pyGinkgo_BUILD_MPI
+       "Build pyGinkgo with MPI/distributed bindings (requires mpi4py)" OFF)
+
+if(pyGinkgo_BUILD_MPI)
+  find_package(MPI REQUIRED COMPONENTS C CXX)
+endif()
 
 # Find pybind11
 find_package(pybind11 QUIET)
@@ -42,10 +58,20 @@ if(NOT nlohmann_json_FOUND)
 endif()
 
 # Find  Ginkgo find_package(Ginkgo 1.7.0 QUIET)
-find_package(ginkgo QUIET)
+find_package(Ginkgo QUIET)
+
+if(pyGinkgo_BUILD_MPI AND Ginkgo_FOUND AND NOT GINKGO_BUILD_MPI)
+  message(
+    FATAL_ERROR
+      "pyGinkgo_BUILD_MPI=ON requires a Ginkgo built with GINKGO_BUILD_MPI=ON. "
+      "The pre-installed Ginkgo at ${Ginkgo_DIR} was built without MPI. "
+      "Either install an MPI-enabled Ginkgo (e.g. the conda 'mpich' or "
+      "'openmpi' variants) or unset Ginkgo_DIR to fetch and build Ginkgo from "
+      "source.")
+endif()
 
 # Fetch Ginkgo if not found
-if(NOT ginkgo_FOUND)
+if(NOT Ginkgo_FOUND)
   # If not found, fetch Ginkgo
   include(FetchContent)
   if(NOT DEFINED GINKGO_BUILD_TESTS)
@@ -57,6 +83,9 @@ if(NOT ginkgo_FOUND)
   if(NOT DEFINED GINKGO_BUILD_EXAMPLES)
     set(GINKGO_BUILD_EXAMPLES OFF)
   endif()
+  if(pyGinkgo_BUILD_MPI)
+    set(GINKGO_BUILD_MPI ON CACHE BOOL "Build Ginkgo with MPI" FORCE)
+  endif()
   FetchContent_Declare(
     ginkgo
     GIT_REPOSITORY https://github.com/ginkgo-project/ginkgo.git
@@ -76,6 +105,37 @@ if(GINKGO_BUILD_SYCL)
   add_compile_definitions(GINKGO_BUILD_SYCL)
 endif()
 
+if(pyGinkgo_BUILD_MPI)
+  list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
+  find_package(Mpi4py REQUIRED)
+
+  # Sanity-check that mpi4py and our MPI agree on a library directory when
+  # mpi4py reports one. Mismatches between mpich/openmpi/nompi flavors here
+  # are a common source of silent ABI breakage downstream.
+  if(Mpi4py_MPI_LIBRARY_DIR)
+    set(_mpi_lib_dirs "")
+    foreach(_lib IN LISTS MPI_C_LIBRARIES)
+      get_filename_component(_dir "${_lib}" DIRECTORY)
+      list(APPEND _mpi_lib_dirs "${_dir}")
+    endforeach()
+    list(REMOVE_DUPLICATES _mpi_lib_dirs)
+    list(FIND _mpi_lib_dirs "${Mpi4py_MPI_LIBRARY_DIR}" _idx)
+    if(_idx EQUAL -1)
+      message(
+        WARNING
+          "mpi4py was built against MPI in '${Mpi4py_MPI_LIBRARY_DIR}' but "
+          "find_package(MPI) selected libraries in '${_mpi_lib_dirs}'. "
+          "This may cause runtime ABI mismatches between mpich/openmpi.")
+    endif()
+  endif()
+
+  include(DetectMpiAbi)
+  pyginkgo_detect_mpi_abi()
+  configure_file(
+    "${CMAKE_CURRENT_SOURCE_DIR}/cmake/pyGinkgo_mpi_abi.hpp.in"
+    "${CMAKE_CURRENT_BINARY_DIR}/generated/pyGinkgo_mpi_abi.hpp" @ONLY)
+endif()
+
 add_subdirectory(src/cpp_bindings)
 
 pybind11_add_module(pyGinkgoBindings ${PYGINKGO_CPP_SOURCES})
@@ -86,6 +146,14 @@ target_link_libraries(pyGinkgoBindings PRIVATE Ginkgo::ginkgo)
 target_link_libraries(pyGinkgoBindings PUBLIC nlohmann_json::nlohmann_json)
 set_property(TARGET pyGinkgoBindings PROPERTY CXX_STANDARD 14)
 
+if(pyGinkgo_BUILD_MPI)
+  target_compile_definitions(pyGinkgoBindings PRIVATE PYGINKGO_BUILD_MPI)
+  target_link_libraries(pyGinkgoBindings PRIVATE MPI::MPI_C MPI::MPI_CXX)
+  target_include_directories(
+    pyGinkgoBindings PRIVATE ${Mpi4py_INCLUDE_DIR}
+                             ${CMAKE_CURRENT_BINARY_DIR}/generated)
+endif()
+
 # Disable false positive warnings from GCC (>= 12.4 and <
 # 13)(https://gcc.gnu.org/bugzilla/show_bug.cgi?id=115824 and
 # https://github.com/pybind/pybind11/issues/5224)
@@ -182,16 +250,20 @@ endif()
 set(CMAKE_INSTALL_RPATH ${Python_SITELIB}/${PROJECT_NAME})
 
 message(STATUS "Install path: ${Python_SITELIB}")
-install(DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}
-        DESTINATION ${Python_SITELIB})
-install(
-  IMPORTED_RUNTIME_ARTIFACTS
-  ginkgo
-  ginkgo_device
-  ginkgo_hip
-  ginkgo_cuda
-  ginkgo_omp
-  ginkgo_dpcpp
-  ginkgo_reference
-  DESTINATION
-  ${Python_SITELIB}/${PROJECT_NAME})
+install(DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}/
+        DESTINATION ${PY_BUILD_CMAKE_MODULE_NAME})
+if(NOT Ginkgo_FOUND)
+  # Only install Ginkgo shared libraries when built from source (FetchContent).
+  # When using a pre-installed Ginkgo (e.g. conda), its libraries are already available.
+  install(
+    IMPORTED_RUNTIME_ARTIFACTS
+    ginkgo
+    ginkgo_device
+    ginkgo_hip
+    ginkgo_cuda
+    ginkgo_omp
+    ginkgo_dpcpp
+    ginkgo_reference
+    DESTINATION
+    ${PY_BUILD_CMAKE_MODULE_NAME})
+endif()

README.md

@@ -242,3 +242,108 @@ x_cupy = cupy.asarray(x_gko)
 ## Benchmarking
 
 The benchmarking results are presented in our [pyGinkgo publication on arXiv](https://arxiv.org/abs/2510.08230).
+
+## Distributed (MPI) bindings
+
+pyGinkgo can optionally expose Ginkgo's distributed-memory types
+(`gko::experimental::distributed::{Partition, Vector, Matrix}`,
+`Schwarz` preconditioner) plus a `PyLinOp` matrix-free trampoline. This
+is gated behind the `pyGinkgo_BUILD_MPI` CMake option so the default
+single-process wheel stays MPI-free.
+
+### Building with MPI
+
+```bash
+cmake -B build -S . -DpyGinkgo_BUILD_MPI=ON
+cmake --build build -j
+pip install ".[mpi]"   # adds the mpi4py runtime dependency
+```
+
+The build will:
+
+1. Refuse to use a pre-installed Ginkgo that was compiled without MPI
+   (the include path must contain `GINKGO_BUILD_MPI=1`).
+2. Locate `mpi4py` via Python and pull in its C headers.
+3. Bake the `MPI_Get_library_version()` string of the build-time MPI
+   into the binary so the Python facade can detect mpich-vs-openmpi
+   mismatches at import time.
+
+### Quick start
+
+```python
+from mpi4py import MPI
+import cupy
+from pyGinkgo import device
+from pyGinkgo.distributed import (
+    Partition, DistributedVector, DistributedMatrix,
+)
+import pyGinkgo.distributed.communicator as gkc
+
+comm = MPI.COMM_WORLD
+exec_ = device("cuda", gkc.map_rank_to_device_id(comm, cupy.cuda.runtime.getDeviceCount()))
+
+# 1D row partition over global_size rows, evenly split across ranks
+part = Partition.uniform(exec_, comm.size, global_size=N)
+
+# Distributed RHS from a local cupy slice
+b = DistributedVector.from_local_array(
+    exec_, comm, global_size=(N, 1), local_array=local_b_cupy)
+```
+
+The C++ `solver.{cg,gmres,bicgstab}_*` factories accept any LinOp, so a
+distributed Matrix or a `PyLinOp` slots into the existing solver path
+unchanged. Each `apply()` returns a `(Convergence, x)` pair so callers
+can read final residual and iteration count:
+
+```python
+cg = pgb.solver.cg_double(exec_, A.raw, max_iters=200,
+                          reduction_factor=1e-10, relative_stop_mode=True)
+log, _ = cg.apply(b.raw, x.raw)
+print(log.get_num_iterations(), log.get_residual_norm())
+```
+
+### Matrix-free distributed operators (`PyLinOp`)
+
+When a `PyLinOp` subclass is used inside a distributed solver, the
+`apply_impl(b, x)` callback receives the *distributed* vectors directly;
+this rank only owns the local block. **The callback is responsible for
+any halo exchange required by the matrix-free matvec** — Ginkgo does not
+perform implicit halo communication for user-defined operators. A
+typical CuPy callback should call `cupy.cuda.Stream.null.synchronize()`
+before returning so Ginkgo (which dispatches on its own executor stream)
+sees the writes.
+
+### Diagnostics: gather a distributed vector to host
+
+```python
+host = v.raw.gather_on_root(0)  # numpy array on rank 0, None elsewhere
+```
+
+### Zero-copy CuPy ↔ distributed.Vector
+
+`DistributedVector.from_local_array(...)` always copies. For a
+zero-copy view that aliases a CuPy buffer, use the C++ static method
+directly and keep a reference to the source array alive:
+
+```python
+v = pgb.distributed.Vector_double.from_local_array_view(
+        exec_, comm, (N, 1), local_cupy_array)
+```
+
+### MPI ABI safety
+
+Mixing an MPICH-built pyGinkgo with an OpenMPI-built mpi4py is a
+common source of mysterious crashes. `pyGinkgo.distributed` performs
+three independent checks:
+
+1. CMake-time: warns if the mpi4py library directory disagrees with the
+   detected MPI install.
+2. Configure-time: a `try_run` probe records the build MPI's
+   `MPI_Get_library_version()` string into the compiled extension.
+3. Import-time: the first call that takes a communicator compares the
+   baked string with `MPI.Get_library_version()` and asks the C++ side
+   to round-trip a small MPI op on the user's `comm`.
+
+A mismatch raises `ImportError` with a hint to install a matching
+conda variant (`pyginkgo-mpi-{mpich,openmpi}`).
+

cmake/DetectMpiAbi.cmake

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: MIT
+#
+# SPDX-FileCopyrightText: 2024 - 2026 pyGinkgo authors
+#
+# Detect the MPI implementation flavor (OpenMPI, MPICH, Intel MPI, ...) and
+# the library version string and bake them into a generated header so we can
+# verify at import time that the mpi4py loaded by the user matches the MPI
+# library pyGinkgo was linked against.
+#
+# Inputs:
+#   MPI_C_COMPILER, MPI_C_INCLUDE_DIRS, MPI_C_LIBRARIES (from find_package(MPI))
+# Outputs (cache vars):
+#   PYGINKGO_MPI_IMPL              - e.g. "OpenMPI", "MPICH", "IntelMPI"
+#   PYGINKGO_MPI_LIBRARY_VERSION   - full string from MPI_Get_library_version
+#
+# A header `pyGinkgo_mpi_abi.hpp` is configured from
+# `pyGinkgo_mpi_abi.hpp.in` into ${CMAKE_CURRENT_BINARY_DIR}/generated/.
+
+function(pyginkgo_detect_mpi_abi)
+  if(DEFINED PYGINKGO_MPI_IMPL AND DEFINED PYGINKGO_MPI_LIBRARY_VERSION)
+    return()
+  endif()
+
+  set(_probe_src "${CMAKE_CURRENT_BINARY_DIR}/pyginkgo_mpi_probe.c")
+  set(_probe_bin "${CMAKE_CURRENT_BINARY_DIR}/pyginkgo_mpi_probe")
+  file(
+    WRITE ${_probe_src}
+    "
+#include <mpi.h>
+#include <stdio.h>
+int main(int argc, char** argv) {
+    char ver[MPI_MAX_LIBRARY_VERSION_STRING];
+    int len = 0;
+    MPI_Init(&argc, &argv);
+    MPI_Get_library_version(ver, &len);
+    fputs(ver, stdout);
+    MPI_Finalize();
+    return 0;
+}
+")
+
+  try_run(
+    _probe_run _probe_compile
+    "${CMAKE_CURRENT_BINARY_DIR}/pyginkgo_mpi_probe_dir"
+    SOURCES "${_probe_src}"
+    LINK_LIBRARIES MPI::MPI_C
+    RUN_OUTPUT_VARIABLE _probe_output
+    COMPILE_OUTPUT_VARIABLE _probe_log)
+
+  if(NOT _probe_compile)
+    message(
+      WARNING
+        "Could not compile MPI ABI probe (cross-compiling?). Falling back to "
+        "MPI_C_VERSION=${MPI_C_VERSION}.\n${_probe_log}")
+    set(_probe_output "Unknown MPI ${MPI_C_VERSION}")
+  endif()
+
+  set(_impl "Unknown")
+  if(_probe_output MATCHES "[Oo]pen MPI" OR _probe_output MATCHES "OpenRTE")
+    set(_impl "OpenMPI")
+  elseif(_probe_output MATCHES "MPICH")
+    set(_impl "MPICH")
+  elseif(_probe_output MATCHES "Intel.* MPI")
+    set(_impl "IntelMPI")
+  endif()
+
+  # Strip newlines / tabs to keep the C string single-line.
+  string(REGEX REPLACE "[\r\n\t]+" " " _probe_output "${_probe_output}")
+  string(STRIP "${_probe_output}" _probe_output)
+
+  set(PYGINKGO_MPI_IMPL
+      "${_impl}"
+      CACHE INTERNAL "Detected MPI implementation")
+  set(PYGINKGO_MPI_LIBRARY_VERSION
+      "${_probe_output}"
+      CACHE INTERNAL "MPI_Get_library_version output")
+  message(
+    STATUS
+      "pyGinkgo MPI: implementation='${PYGINKGO_MPI_IMPL}', version='${PYGINKGO_MPI_LIBRARY_VERSION}'"
+  )
+endfunction()