Add GTL

docs/src/configuration.md

@@ -60,7 +60,6 @@ If the implementation is changed, you will need to call this function again. See
     from transitive dependencies is broken ([Preferences.jl#24](https://github.com/JuliaPackaging/Preferences.jl/issues/24)).
     To fix this update your version of Julia, or add `MPIPreferences` as a direct dependency to your project.
 
-
 ### Notes to HPC cluster administrators
 
 Preferences are merged across the Julia load path, such that it is feasible to provide a module file that appends a path to
@@ -107,6 +106,50 @@ Preferences are merged across the Julia load path, such that it is feasible to p
    that will take precedent by modifying the local `Project.toml` or by providing a
    `LocalPreferences.toml` file.
 
+### Notes about vendor-provided MPI backends
+
+`MPIPreferences` can load vendor-specific libraries and settings using the
+`vendor` parameter, eg `MPIPreferences.use_system_binary(mpiexec="srun", vendor="cray")`
+configures `MPIPreferences` for use on Cray systems with `srun`.
+
+!!! note
+    Currently `vendor` only supports Cray systems.
+
+This populates the `library_names`, `preloads`, `preloads_env_switch` and
+`cclibs` preferences. These are defermined by parsing `cc --cray-print-opts=all`
+emitted from the Cray Compiler Wrappers. Therefore `use_system_binary` needs
+to be run on the target system, with the corresponding `PrgEnv` loaded.
+
+The function of these settings are as follows:
+ * `preloads` specifies a list of libraries that are to be loaded (in order)
+   before `libmpi`.
+ * `preloads_env_switch` specifies the name of an environment variable that, if
+   set to `0`, can disable the `preloads`
+ * `cclibs` is a list of libraries also linked by the compiler wrappers. This is
+   recorded mainly for debugging purposes, and the libraries listed here are not
+   explicitly loaded by `MPI.jl`.
+
+If these are set, the `_format` key will be set to `"1.1"`.
+
+An example of running `MPIPreferences.use_system_library(vendor="cray")` in
+`PrgEnv-gnu` is:
+
+```toml
+[MPIPreferences]
+_format = "1.1"
+abi = "MPICH"
+binary = "system"
+cclibs = ["cupti", "cudart", "cuda", "sci_gnu_82_mpi", "sci_gnu_82", "dl", "dsmml", "xpmem"]
+libmpi = "libmpi_gnu_91.so"
+mpiexec = "mpiexec"
+preloads = ["libmpi_gtl_cuda.so"]
+preloads_env_switch = "MPICH_GPU_SUPPORT_ENABLED"
+```
+
+This is an example of CrayMPICH requiring `libmpi_gtl_cuda.so` to be preloaded,
+unless `MPICH_GPU_SUPPORT_ENABLED=0` (the latter allowing MPI-enabled code to
+run on a non-GPU enabled node without needing a seperate `LocalPreferences.toml`).
+
 ## [Using an alternative JLL-provided MPI library](@id configure_jll_binary)
 
 The following MPI implementations are provided as JLL packages and automatically obtained when installing MPI.jl:

lib/MPIPreferences/src/MPIPreferences.jl

@@ -4,7 +4,7 @@ export use_jll_binary, use_system_binary
 
 using Preferences, Libdl
 
-if !(VersionNumber(@load_preference("_format", "1.0")) <= v"1.0")
+if !(VersionNumber(@load_preference("_format", "1.1")) <= v"1.1")
     error("The preferences attached to MPIPreferences are incompatible with this version of the package.")
 end
 
@@ -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,13 @@ Options:
   using [`identify_abi`](@ref). See [`abi`](@ref) for currently supported
   values.
 
+- `vendor`: can be either `nothing` or a vendor name (such a `"cray"`). If
+  `vendor` has the value "cray", then the output from `cc --cray-print-opts=all`
+  is parsed for which libraries are linked by the Cray Compiler Wrappers. Note
+  that if `mpi_gtl_*` is present, then this .so will be added to the preloads.
+  Also note that the inputs to `library_names` will be overwritten by the
+  library name used by the compiler wrapper.
+
 - `export_prefs`: if `true`, the preferences into the `Project.toml` instead of
   `LocalPreferences.toml`.
 
@@ -145,10 +157,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 +185,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",
+        "_format" => isnothing(vendor) ? "1.0" : "1.1",
         "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 +216,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,66 @@
+module CrayParser
+
+filter(f::Function)::Function = Base.Fix1(Base.filter, f)
+map(f::Function)::Function    = Base.Fix1(Base.map, f)
+reduce(f::Function)::Function = Base.Fix1(Base.reduce, f)
+
+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
+
+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))
+    return only(x)
+end
+
+function cray_gtl(libs)
+    x = libs |>
+        filter(x->startswith(x, libmpi_prefix)) |> 
+        filter(x->startswith(x, libgtl_prefix))
+    return only(x)
+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 = readchomp(Cmd(["cc", "--cray-print-opts=all"]))
+
+    ld_paths = SubString{String}[]
+    libs = SubString{String}[]
+    for opt in split(cray_opts, " ")        |>
+               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

@@ -113,6 +113,16 @@ 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
+        if !isnothing(preloads)
+            if isnothing(preloads_env_switch) || get(ENV, preloads_env_switch, "0") == "1"
+                for preload in preloads
+                    Libdl.dlopen(preload, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
+                end
+            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: MPICH_jll, libmpi, libmpi_handle, mpiexec
     const libmpiconstants = nothing
+    const preloads = []
+    const preloads_env_switch = nothing
 elseif MPIPreferences.binary == "OpenMPI_jll"
     import OpenMPI_jll: OpenMPI_jll, libmpi, libmpi_handle, mpiexec
     const libmpiconstants = nothing
+    const preloads = []
+    const preloads_env_switch = nothing
 elseif MPIPreferences.binary == "MicrosoftMPI_jll"
     import MicrosoftMPI_jll: 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)")