Add GTL

lib/MPIPreferences/src/MPIPreferences.jl

@@ -50,6 +50,7 @@ else
     error("Unknown binary: $binary")
 end
 
+include("parse_cray_cc.jl")
 @static if binary == "system"
     include("system.jl")
 end
@@ -81,7 +82,10 @@ function use_jll_binary(binary = Sys.iswindows() ? "MicrosoftMPI_jll" : "MPICH_j
         "binary" => binary,
         "libmpi" => nothing,
         "abi" => nothing,
-        "mpiexec" => nothing;
+        "mpiexec" => nothing,
+        "preloads" => [],
+        "preloads_env_switch" => nothing,
+        "cclibs" => nothing;
         export_prefs=export_prefs,
         force=force
     )
@@ -114,6 +118,7 @@ end
         mpiexec = "mpiexec",
         abi = nothing,
         export_prefs = false,
+        gtl_names=nothing,
         force = true)
 
 Switches the underlying MPI implementation to a system provided one. A restart
@@ -136,6 +141,8 @@ Options:
   using [`identify_abi`](@ref). See [`abi`](@ref) for currently supported
   values.
 
+- `vendor`: can be either `nothing` or a vendor name (such a `"cray"`).
+
 - `export_prefs`: if `true`, the preferences into the `Project.toml` instead of
   `LocalPreferences.toml`.
 
@@ -145,10 +152,23 @@ function use_system_binary(;
         library_names=["libmpi", "libmpi_ibm", "msmpi", "libmpich", "libmpi_cray", "libmpitrampoline"],
         mpiexec="mpiexec",
         abi=nothing,
+        vendor=nothing,
         export_prefs=false,
-        force=true,
+        force=true
     )
     binary = "system"
+    # vendor workarounds
+    preloads = []
+    preloads_env_switch = nothing
+    cclibs = []
+    if vendor == "cray"
+        cray_pe = CrayParser.analyze_cray_cc()
+        library_names = [cray_pe.libmpi]
+        preloads = [cray_pe.libgtl]
+        preloads_env_switch = cray_pe.gtl_env_switch
+        cclibs = cray_pe.cclibs
+    end
+
     # Set `ZES_ENABLE_SYSMAN` to work around https://github.com/open-mpi/ompi/issues/10142
     libmpi = withenv("ZES_ENABLE_SYSMAN" => "1") do
         find_library(library_names)
@@ -160,23 +180,28 @@ function use_system_binary(;
             If you want to try different name(s) for the MPI library, use
                 MPIPreferences.use_system_binary(; library_names=[...])""")
     end
+
     if isnothing(abi)
         abi = identify_abi(libmpi)
     end
+
     if mpiexec isa Cmd
         mpiexec = collect(mpiexec)
     end
+
     set_preferences!(MPIPreferences,
         "_format" => "1.0",
         "binary" => binary,
         "libmpi" => libmpi,
         "abi" => abi,
         "mpiexec" => mpiexec,
+        "preloads" => preloads,
+        "preloads_env_switch" => preloads_env_switch,
+        "cclibs" => cclibs;
         export_prefs=export_prefs,
         force=force
     )
 
-
     if VERSION <= v"1.6.5" || VERSION == v"1.7.0"
         @warn """
         Due to a bug in Julia (until 1.6.5 and 1.7.1), setting preferences in transitive dependencies
@@ -186,10 +211,10 @@ function use_system_binary(;
     end
 
     if binary == MPIPreferences.binary && abi == MPIPreferences.abi && libmpi == System.libmpi && mpiexec == System.mpiexec_path
-        @info "MPIPreferences unchanged" binary libmpi abi mpiexec
+        @info "MPIPreferences unchanged" binary libmpi abi mpiexec preloads preloads_env_switch
     else
         PREFS_CHANGED[] = true
-        @info "MPIPreferences changed" binary libmpi abi mpiexec
+        @info "MPIPreferences changed" binary libmpi abi mpiexec preloads preloads_env_switch
 
         if DEPS_LOADED[]
             error("You will need to restart Julia for the changes to take effect")

lib/MPIPreferences/src/parse_cray_cc.jl

@@ -0,0 +1,94 @@
+module CrayParser
+
+import Base: filter, map, reduce
+filter(f::Function)::Function = x -> filter(f, x)
+map(f::Function)::Function    = x -> map(f, x)
+reduce(f::Function)::Function = x -> reduce(f, x)
+
+struct CrayPE
+    libmpi::String
+    libgtl::String
+    cclibs::Vector{String}
+    gtl_env_switch::String
+
+    CrayPE(mpi_dl::T, gtl_dl::T, cclibs::Vector{T}) where T <:AbstractString = new(
+        "lib" * mpi_dl * ".so", # Assuming Linux -- CrayPE is only avaialbe for linux anyway
+        "lib" * gtl_dl * ".so",
+        cclibs,
+        "MPICH_GPU_SUPPORT_ENABLED"
+    )
+end
+
+function communicate(cmd::Cmd, input="")
+    inp = Pipe()
+    out = Pipe()
+    err = Pipe()
+
+    process = run(pipeline(cmd, stdin=inp, stdout=out, stderr=err), wait=false)
+
+    close(out.in)
+    close(err.in)
+
+    stdout = @async String(read(out))
+    stderr = @async String(read(err))
+
+    write(process, input)
+    close(inp)
+
+    wait(process)
+    return (
+        stdout = fetch(stdout),
+        stderr = fetch(stderr),
+        code = process.exitcode
+    )
+end
+
+const libmpi_prefix = "mpi_"
+const libgtl_prefix = "mpi_gtl_"
+
+function cray_mpi(libs)
+    x = libs |>
+        filter(x-> startswith(x, libmpi_prefix)) |> 
+        filter(x->!startswith(x, libgtl_prefix))
+    @assert length(x) == 1
+    return x[1]
+end
+
+function cray_gtl(libs)
+    x = libs |>
+        filter(x->startswith(x, libmpi_prefix)) |> 
+        filter(x->startswith(x, libgtl_prefix))
+    @assert length(x) == 1
+    return x[1]
+end
+
+function other_libs(libs)
+    x = libs |>
+        filter(x->!startswith(x, libmpi_prefix)) |> 
+        filter(x->!startswith(x, libgtl_prefix))
+    return x
+end
+
+function analyze_cray_cc()
+    cray_opts = communicate(Cmd(["cc", "--cray-print-opts=all"]))
+    @assert cray_opts.code == 0
+
+    ld_paths = SubString{String}[]
+    libs = SubString{String}[]
+    for opt in split(cray_opts.stdout, " ") |>
+               map(x->split(x, ","))        |>
+               reduce(vcat)                 |>
+               map(x->replace(x, "\n"=>""))
+        if startswith(opt, "-L")
+            push!(ld_paths, @view opt[3:end])
+        end
+
+        if startswith(opt, "-l")
+            push!(libs, @view opt[3:end])
+        end
+    end
+
+    CrayPE(cray_mpi(libs), cray_gtl(libs), other_libs(libs))
+end
+
+end

lib/MPIPreferences/src/system.jl

@@ -2,12 +2,37 @@ module System
     export libmpi, mpiexec
     using Preferences, Libdl
     const libmpi = @load_preference("libmpi")
+    const preloads = @load_preference("preloads")
+    const preloads_env_switch = @load_preference("preloads_env_switch")
     const mpiexec_path = @load_preference("mpiexec")
     mpiexec(;adjust_PATH=true, adjust_LIBPATH=true) = `$mpiexec_path`
     mpiexec(f;adjust_PATH=true, adjust_LIBPATH=true) = f(`$mpiexec_path`)
 
     libmpi_handle = C_NULL
     function __init__()
+        # preload any dependencies of libmpi (if needed, eg. GTL on cray) before
+        # dlopen'ing the MPI library: https://github.com/JuliaParallel/MPI.jl/pull/716
+        preload_enabled = false
+        if isnothing(preloads_env_switch)
+            preload_enabled = true
+        elseif get(ENV, preloads_env_switch, "0") == "1"
+            preload_enabled = true
+        end
+
+        if preload_enabled
+            for preload in preloads
+                try
+                    Libdl.dlopen(preload, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
+                catch error
+                    @error """
+                        $(preload) could not be loaded, see error message below.
+                        Use `MPIPreferences.use_system_binary` or `MPIPreferences.use_jll_binary` to reconfigure the package and then restart Julia.
+                        """ error
+                end
+                # TODO: do we want to expose the preload handles?
+            end
+        end
+
         global libmpi_handle = try
             Libdl.dlopen(libmpi, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
         catch error

src/MPI.jl

@@ -103,6 +103,20 @@ function __init__()
     end
 
     @static if Sys.isunix()
+        # preload any dependencies of libmpi (if needed, eg. GTL on cray) before
+        # dlopen'ing the MPI library: https://github.com/JuliaParallel/MPI.jl/pull/716
+        preload_enabled = false
+        if isnothing(preloads_env_switch)
+            preload_enabled = true
+        elseif get(ENV, preloads_env_switch, "0") == "1"
+            preload_enabled = true
+        end
+        if preload_enabled
+            for preload in preloads
+                Libdl.dlopen(preload, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
+            end
+        end
+
         # dlopen the MPI library before any ccall:
         # - RTLD_GLOBAL is required for Open MPI
         #   https://www.open-mpi.org/community/lists/users/2010/04/12803.php

src/api/api.jl

@@ -3,25 +3,39 @@ module API
 export MPI_Aint, MPI_Count, MPI_Offset, MPI_Status,
     MPI_Comm, MPI_Datatype, MPI_Errhandler, MPI_File, MPI_Group,
     MPI_Info, MPI_Message, MPI_Op, MPI_Request, MPI_Win,
-    libmpi, mpiexec, @mpichk, @mpicall, MPIPtr, SentinelPtr, FeatureLevelError
+    libmpi, mpiexec, @mpichk, @mpicall, MPIPtr, SentinelPtr, FeatureLevelError,
+    preloads, preloads_env_switch
+    # libgtl
 
 import MPIPreferences
 using Libdl
 
 if MPIPreferences.binary == "MPICH_jll"
     import MPICH_jll: libmpi, libmpi_handle, mpiexec
     const libmpiconstants = nothing
+    const preloads = []
+    const preloads_env_switch = nothing
 elseif MPIPreferences.binary == "OpenMPI_jll"
     import OpenMPI_jll: libmpi, libmpi_handle, mpiexec
     const libmpiconstants = nothing
+    const preloads = []
+    const preloads_env_switch = nothing
 elseif MPIPreferences.binary == "MicrosoftMPI_jll"
     import MicrosoftMPI_jll: libmpi, libmpi_handle, mpiexec
     const libmpiconstants = nothing
+    const preloads = []
+    const preloads_env_switch = nothing
 elseif MPIPreferences.binary == "MPItrampoline_jll"
     import MPItrampoline_jll: MPItrampoline_jll, libmpi, libmpi_handle, mpiexec
     const libmpiconstants = MPItrampoline_jll.libload_time_mpi_constants_path
+    # TODO: We'll probably need the preloads -- like cray's GTL -- with
+    # libmpitrampoline, and until MPItrampoline_jll "understands" preloads, this
+    # should be an acceptable workaround
+    const preloads = MPIPreferences.Preferences.@load_preference("preloads")
+    const preloads_env_switch = MPIPreferences.Preferences.@load_preference("preloads_env_switch")
 elseif MPIPreferences.binary == "system"
-    import MPIPreferences.System: libmpi, libmpi_handle, mpiexec
+    import MPIPreferences.System: libmpi, libmpi_handle, mpiexec,
+        preloads, preloads_env_switch
     const libmpiconstants = nothing
 else
     error("Unknown MPI binary: $(MPIPreferences.binary)")