Add GTL (#716)

Support for Cray MPI GPU-aware

Co-authored-by: Simon Byrne <[email protected]>
Co-authored-by: Valentin Churavy <[email protected]>

Project.toml

@@ -22,7 +22,7 @@ Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 AMDGPU = "0.3, 0.4, 0.5"
 CUDA = "3, 4"
 DocStringExtensions = "0.8, 0.9"
-MPIPreferences = "0.1.6"
+MPIPreferences = "0.1.8"
 PkgVersion = "0.3"
 PrecompileTools = "1.0.1"
 Requires = "~0.5, 1.0"

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:

docs/src/reference/mpipreferences.md

@@ -22,6 +22,7 @@ MPIPreferences.use_jll_binary
 ```@docs
 MPIPreferences.check_unchanged
 MPIPreferences.identify_abi
+MPIPreferences.dlopen_preloads
 ```
 
 ## Preferences schema

lib/MPIPreferences/Project.toml

@@ -1,7 +1,7 @@
 name = "MPIPreferences"
 uuid = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"
 authors = []
-version = "0.1.8"
+version = "0.1.9"
 
 [deps]
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

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.0")) <= v"1.1")
     error("The preferences attached to MPIPreferences are incompatible with this version of the package.")
 end
 
@@ -50,6 +50,11 @@ else
     error("Unknown binary: $binary")
 end
 
+include("preloads.jl")
+using .Preloads: dlopen_preloads, preloads, preloads_env_switch
+
+include("parse_cray_cc.jl")
+
 @static if binary == "system"
     include("system.jl")
 end
@@ -81,7 +86,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
     )
@@ -113,6 +121,7 @@ end
         library_names = ["libmpi", "libmpi_ibm", "msmpi", "libmpich", "libmpi_cray", "libmpitrampoline"],
         mpiexec = "mpiexec",
         abi = nothing,
+        vendor = nothing,
         export_prefs = false,
         force = true)
 
@@ -136,6 +145,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 +161,26 @@ 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 === nothing
+    elseif 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
+    else
+        error("Unknown vendor $vendor")
+    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 +192,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 +223,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/preloads.jl

@@ -0,0 +1,44 @@
+module Preloads
+
+any(f::Function)::Function = Base.Fix1(Base.any, f)
+
+using Preferences, Libdl
+
+const preloads = @load_preference("preloads")
+const preloads_env_switch = @load_preference("preloads_env_switch")
+
+
+function is_loaded(name)
+    if dllist() |> any(x->endswith(x, name))
+        return true
+    end
+    return false
+end
+
+"""
+    dlopen_preloads()
+
+dlopen's all preloads specified in the preloads section of MPIPreferences
+"""
+function dlopen_preloads()
+    if !isnothing(preloads)
+        if isnothing(preloads_env_switch) || get(ENV, preloads_env_switch, "0") == "1"
+            for preload in preloads
+                if is_loaded(preload)
+                    continue
+                end
+
+                try
+                    dlopen(preload, RTLD_LAZY | RTLD_GLOBAL)
+                catch error
+                    @error """
+                        $(preload) could not be loaded, see error message below.
+                        Use `MPIPreferences` to reconfigure the package and then restart Julia.
+                        """ error
+                end
+            end
+        end
+    end
+end
+
+end

lib/MPIPreferences/src/system.jl

@@ -1,13 +1,21 @@
 module System
+    include("preloads.jl")
+
     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
+        Preloads.dlopen_preloads()
+
         global libmpi_handle = try
             Libdl.dlopen(libmpi, Libdl.RTLD_LAZY | Libdl.RTLD_GLOBAL)
         catch error

src/MPI.jl

@@ -112,6 +112,10 @@ function __init__()
         """ ENV["JULIA_MPI_BINARY"]=mpi_env_binary MPIPreferences.binary
     end
 
+    # 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
+    MPIPreferences.dlopen_preloads()
+
     @static if Sys.isunix()
         # dlopen the MPI library before any ccall:
         # - RTLD_GLOBAL is required for Open MPI