Add gc_safe option to @mpicall and @mpichk macros

.github/workflows/Documenter.yml

@@ -42,3 +42,9 @@ jobs:
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         run: julia --color=yes --project=docs/ docs/make.jl
+      - name: "Upload docs build as artifact"
+        uses: actions/upload-artifact@v7
+        with:
+          name: documentation-build
+          path: docs/build
+          retention-days: 30

.gitignore

@@ -1,4 +1,5 @@
 *.cov
 *.mem
 Manifest.toml
+Manifest-v*.toml
 LocalPreferences.toml

docs/src/knownissues.md

@@ -143,6 +143,21 @@ export UCX_ERROR_SIGNALS="SIGILL,SIGBUS,SIGFPE"
 ```
 before calling `mpiexec`.
 
+## Multi-threading and garbage collection
+
+On Julia v1.12 and later, blocking MPI calls (e.g. `MPI_Wait`, `MPI_Barrier`, blocking point-to-point communications and collectives) are marked as *GC-safe*: the Julia garbage collector running on one thread does not have to wait for other threads which are inside one of these MPI calls.
+This prevents the whole process from stalling — or deadlocking outright — when a thread blocks in MPI while waiting on the progress of other ranks and another thread needs to run a garbage collection.
+
+A consequence of this is that in a multi-threaded program the garbage collector can now run *concurrently* with a blocking MPI call, and its finalizers can free MPI handles (communicators, requests, datatypes, operators, etc.) by calling functions like `MPI_Comm_free` while another thread is still inside MPI.
+Concurrent MPI calls from different threads are only permitted when MPI is initialized with the `MPI.THREAD_MULTIPLE` [`MPI.ThreadLevel`](@ref), whereas [`MPI.Init`](@ref) defaults to `:serialized`.
+Multi-threaded programs should therefore initialize MPI with
+
+```julia
+MPI.Init(; threadlevel=:multiple)
+```
+
+Single-threaded programs are unaffected.
+
 ## CUDA-aware MPI
 
 ### Memory pool

docs/src/reference/api.md

@@ -3,6 +3,11 @@
 The `MPI.API` submodule provides a low-level interface which closely matches the MPI C API.
 While these functions are not intended for general usage, they are useful for calling MPI routines not yet available in `MPI.jl` main interface, and is the basis for the high-level wrappers. The methods suffixed with `_c` allow `MPI_count` typed arguments (vs `int` for the standard ones). The size of `MPI_count` depends on the implementation, but usually allows `64bit` integer offsets.
 
+```@docs
+MPI.API.@mpicall
+MPI.API.@mpichk
+```
+
 ```@autodocs
 Modules = [MPI.API]
 Order = [:function]

gen/Project.toml

@@ -8,6 +8,6 @@ MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 OpenMPI_jll = "fe0851c0-eecd-5654-98d4-656369965a5c"
 
 [compat]
-Clang = "0.17.1"
+Clang = "0.19"
 MPIPreferences = "0.1.3"
-julia = "1.6"
+julia = "1.12"

gen/src/MPIgenerator.jl

@@ -18,7 +18,9 @@ module MPIgenerator
         mkpath(out)
 
         options = load_options(joinpath(@__DIR__, "generator.toml"))  # wrapper generator options
-        options["general"]["callback_documentation"] = node -> [string('$', "(_doc_external(:", node.id, "))")]
+        # NOTE: in Clang.jl v0.19+ the callback takes also the docstring extracted from the
+        # C comments as second argument, which we ignore
+        options["general"]["callback_documentation"] = (node, doc) -> [string('$', "(_doc_external(:", node.id, "))")]
 
         include_dir = normpath(artifact_dir, "include")
 
@@ -48,6 +50,52 @@ module MPIgenerator
             :MPI_Wtick,
         )
 
+        # these methods are called with `gc_safe=true`, to allow the garbage
+        # collector to run concurrently with the (potentially blocking) MPI call.
+        # The criterion for inclusion is that the worst-case duration of the call
+        # is unbounded because it waits on the progress of peer ranks, rather than
+        # being bounded by local work, and that the call can only re-enter Julia
+        # through `@cfunction` callbacks (e.g. custom reduction operators), which
+        # are safe in `gc_safe` regions
+        gc_safe = (
+            # request completion
+            :MPI_Wait,
+            :MPI_Waitall,
+            :MPI_Waitany,
+            :MPI_Waitsome,
+            :MPI_Mrecv,
+            # blocking point-to-point
+            :MPI_Send,
+            :MPI_Ssend,
+            :MPI_Recv,
+            :MPI_Sendrecv,
+            :MPI_Sendrecv_replace,
+            # blocking collectives
+            :MPI_Barrier,
+            :MPI_Bcast,
+            :MPI_Gather,
+            :MPI_Gatherv,
+            :MPI_Scatter,
+            :MPI_Scatterv,
+            :MPI_Allgather,
+            :MPI_Allgatherv,
+            :MPI_Alltoall,
+            :MPI_Alltoallv,
+            :MPI_Alltoallw,
+            :MPI_Reduce,
+            :MPI_Allreduce,
+            :MPI_Reduce_scatter,
+            :MPI_Reduce_scatter_block,
+            :MPI_Scan,
+            :MPI_Exscan,
+            # blocking neighborhood collectives
+            :MPI_Neighbor_allgather,
+            :MPI_Neighbor_allgatherv,
+            :MPI_Neighbor_alltoall,
+            :MPI_Neighbor_alltoallv,
+            :MPI_Neighbor_alltoallw,
+        )
+
         versioned = Dict(
             :MPI_Dist_graph_create_adjacent => v"2.2",
             :MPI_Dist_graph_neighbors_count => v"2.2",
@@ -59,19 +107,45 @@ module MPIgenerator
             :MPI_Neighbor_alltoall => v"3.0",
         )
 
-        src, fn = joinpath(out, "api.jl"), replace(@__FILE__, r".*MPI.jl" => "MPI.jl")
-        lines = String["# WARNING: this signature file for $(MPIPreferences.binary) has been auto-generated, please edit $fn instead !\n"]
+        src, fn = joinpath(out, "api.jl"), joinpath("MPI.jl", relpath(@__FILE__, normpath(@__DIR__, "..", "..")))
+        lines = String[]
         for line in readlines(src)
             if (m = match(r"^ccall.*:([\w_]+)", lstrip(line))) ≢ nothing
                 sym = first(m.captures) |> Symbol
-                repl = sym ∈ mpicall ? "@mpicall ccall" : "@mpichk ccall"
+                repl = (sym ∈ mpicall ? "@mpicall" : "@mpichk") *
+                    (sym ∈ gc_safe ? " gc_safe=true" : "") * " ccall"
                 line = replace(line, "Ptr{Cvoid}" => "MPIPtr", "ccall" => repl)
                 if (ver = get(versioned, sym, nothing)) ≢ nothing
                     line *= " $(repr(ver))"
                 end
             end
             push!(lines, replace(line, raw"\$" => '$'))
         end
+
+        # group the lines into blocks (docstring + function definition) and sort them by
+        # name of the function, to make the output stable: the order in which the
+        # functions are declared in the headers may change between releases of the MPI
+        # implementations
+        blocks = Vector{Vector{String}}()
+        block = String[]
+        for line in lines
+            # drop the blank lines separating the blocks, they are re-added below
+            isempty(block) && isempty(line) && continue
+            push!(block, line)
+            if line == "end"  # end of a function definition
+                push!(blocks, block)
+                block = String[]
+            end
+        end
+        isempty(block) || error("malformed generated file $src, trailing lines: $block")
+        function_name(block) = match(r"^function (\w+)", block[findfirst(startswith("function "), block)]).captures[1]
+        sort!(blocks; by=function_name)
+
+        lines = String["# WARNING: this signature file for $(MPIPreferences.binary) has been auto-generated, please edit $fn instead !\n"]
+        for block in blocks
+            append!(lines, block)
+            push!(lines, "")
+        end
         write(src, join(lines, "\n"))
 
         dst = normpath(@__DIR__, "..", "..", "src", "api", "generated_api.jl")

src/MPI.jl

@@ -21,7 +21,7 @@ end
 function _doc_external(fname)
     # Ideally we'd like to use the "latest" version of the docs, but MPICH v4.0
     # at the moment seems to be the last version with all the docstrings.
-    mpich = "[MPICH](https://www.mpich.org/static/docs/v4.0/www3/$(fname).html)"
+    mpich = "[MPICH](https://www.mpich.org/static/docs/v5.0.1/www3/$(fname).html)"
     # All the *_c functions are undocumented in OpenMPI website
     if !endswith(string(fname), "_c")
         openmpi = "[OpenMPI](https://docs.open-mpi.org/en/main/man-openmpi/man3/$(fname).3.html)"

src/api/api.jl

@@ -89,7 +89,19 @@ end
 
 const use_stdcall = startswith(basename(libmpi), "msmpi")
 
-macro mpicall(expr)
+# Parse an optional leading `gc_safe=(true|false)` argument (as in `@ccall`).
+# Returns (gc_safe::Bool, remaining_args).
+function parse_gc_safe(args)
+    if !isempty(args) && Meta.isexpr(args[1], :(=), 2) && args[1].args[1] === :gc_safe
+        val = args[1].args[2]
+        val === true || val === false ||
+            throw(ArgumentError("`gc_safe` must be literally `true` or `false`"))
+        return val, args[2:end]
+    end
+    return false, args
+end
+
+function mpicall_lower(expr, gc_safe::Bool)
     @assert expr isa Expr && expr.head == :call && expr.args[1] == :ccall
 
     # On unix systems we call the global symbols to allow for LD_PRELOAD interception
@@ -103,10 +115,37 @@ macro mpicall(expr)
     # this only affects 32-bit Windows
     # unfortunately we need to use ccall to call Get_library_version
     # so check using library name instead
-    if use_stdcall
+    if gc_safe && VERSION >= v"1.12"
+        # same lowering as Base's `@ccall gc_safe=true` (`ccall_macro_lower` in
+        # `base/c.jl`); the `gc_safe` flag is only supported by Julia 1.12+
+        convention = use_stdcall ? :stdcall : :ccall
+        insert!(expr.args, 3, Expr(:cconv, (convention, UInt16(0), true), 0))
+    elseif use_stdcall
         insert!(expr.args, 3, :stdcall)
     end
-    return esc(expr)
+    return expr
+end
+
+"""
+    @mpicall [gc_safe=false] ccall((:MPI_Fn, libmpi), rettype, (argtypes...), args...)
+
+Wrapper around `ccall` for calling MPI library functions, handling the
+platform-specific calling convention (`stdcall` for Microsoft MPI on 32-bit
+Windows) and symbol lookup.
+
+Setting `gc_safe=true` allows the garbage collector to run concurrently with
+the call, which is useful for MPI functions which may block.  It is ignored on
+Julia versions older than v1.12, which don't support this option.
+
+!!! warning
+    `gc_safe=true` can lead to undefined behavior if the MPI function calls
+    back into the Julia runtime, see the documentation of `@ccall`.
+"""
+macro mpicall(args...)
+    gc_safe, args = parse_gc_safe(args)
+    length(args) == 1 ||
+        throw(ArgumentError("@mpicall takes a `ccall(...)` expression and an optional leading `gc_safe=(true|false)` argument"))
+    esc(mpicall_lower(args[1], gc_safe))
 end
 
 """
@@ -122,7 +161,22 @@ function Base.show(io::IO, err::FeatureLevelError)
     print(io, "FeatureLevelError($(err.function_name)): Minimum MPI version is $(err.min_version)")
 end
 
-macro mpichk(expr, min_version=nothing)
+"""
+    @mpichk [gc_safe=false] ccall((:MPI_Fn, libmpi), Cint, (argtypes...), args...) [min_version]
+
+Like [`@mpicall`](@ref), but checks the returned error code and throws an
+[`MPIError`](@ref) if the call was not successful.
+
+If the minimal MPI version `min_version` required for `MPI_Fn` to be available
+is provided and the function is not found in the MPI library, a
+[`FeatureLevelError`](@ref) is thrown instead of performing the call.
+"""
+macro mpichk(args...)
+    gc_safe, args = parse_gc_safe(args)
+    1 <= length(args) <= 2 ||
+        throw(ArgumentError("@mpichk takes a `ccall(...)` expression, an optional leading `gc_safe=(true|false)` argument, and an optional trailing minimum MPI version"))
+    expr = args[1]
+    min_version = length(args) == 2 ? args[2] : nothing
     if !isnothing(min_version) && expr.args[2].head == :tuple
         fn = expr.args[2].args[1].value
         if isnothing(dlsym(libmpi_handle, fn; throw_error=false))
@@ -132,7 +186,7 @@ macro mpichk(expr, min_version=nothing)
         end
     end
 
-    expr = macroexpand(@__MODULE__, :(@mpicall($expr)))
+    expr = mpicall_lower(expr, gc_safe)
     # MPI_SUCCESS is defined to be 0
     :((errcode = $(esc(expr))) == 0 || throw(MPIError(errcode)))
 end