feat: multi GPU support (#587)

* feat: create PJRT client on all GPUs

* fix: remove return

* fix: indexing

* fix: cuda_visible_devices handling

* revert: to old state

* feat: add devices api

* feat: expose more build options

* fix: remove function

* fix: build_local.jl

* feat: correct device handling for concreterarray

* docs: add to docs

* feat: correctly set build options

* fix: empty device list

* fix: precompilation???

* Update WORKSPACE

* chore: bump version

* fix: precompilation

---------

Co-authored-by: William Moses <[email protected]>

Author: avik-pal

Project.toml

@@ -70,7 +70,7 @@ PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.4"
-Reactant_jll = "0.0.51"
+Reactant_jll = "0.0.52"
 Scratch = "1.2"
 Sockets = "1.10"
 SpecialFunctions = "2.4"

deps/Project.toml

@@ -1,10 +1,11 @@
 [deps]
+ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
+BinaryBuilderBase = "7f725544-6523-48cd-82d1-3fa08ff4056e"
+Clang = "40e3b903-d033-50b4-a0cc-940c62c95e31"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"
-Clang = "40e3b903-d033-50b4-a0cc-940c62c95e31"
-BinaryBuilderBase = "7f725544-6523-48cd-82d1-3fa08ff4056e"
 
 [compat]
 Clang = "0.18"

deps/ReactantExtra/API.cpp

@@ -388,6 +388,10 @@ extern "C" PjRtDevice *ClientGetAddressableDevice(PjRtClient *client,
       client->LookupAddressableDevice(PjRtLocalDeviceId(device_id)));
 }
 
+extern "C" const char *ClientGetPlatformName(PjRtClient *client) {
+  return cstr_from_string(client->platform_name());
+}
+
 // To keep in sync with JLAllocatorStats in src/XLA.jl
 struct JLAllocatorStats {
   int64_t num_allocs;
@@ -578,14 +582,22 @@ extern "C" MlirModule ConvertLLVMStrToMLIR(const char *lmod, MlirContext cctx) {
 
 /* Note that this */
 extern "C" xla::PjRtLoadedExecutable *ClientCompile(PjRtClient *client,
-                                                    MlirModule cmod) {
+                                                    MlirModule cmod,
+                                                    int device_ordinal,
+                                                    int num_replicas,
+                                                    int num_partitions,
+                                                    bool use_shardy_partitioner) {
   auto program =
       std::make_unique<xla::ifrt::HloProgram>(cast<ModuleOp>(*unwrap(cmod)));
 
   CompileOptions options;
-  // options.argument_layouts;
-  // options.executable_build_options.set_device_ordinal();
-  // options.executable_build_options.set_result_layout();
+
+  if (device_ordinal >= 0) {
+    options.executable_build_options.set_device_ordinal(device_ordinal);
+  }
+  options.executable_build_options.set_num_replicas(num_replicas);
+  options.executable_build_options.set_num_partitions(num_partitions);
+  options.executable_build_options.set_use_shardy_partitioner(use_shardy_partitioner);
 
   auto addressable_devices = client->addressable_devices();
   if (!addressable_devices.empty()) {

deps/ReactantExtra/WORKSPACE

@@ -9,7 +9,7 @@ http_archive(
     urls = ["https://github.com/wsmoses/nsync/archive/{commit}.tar.gz".format(commit = NSYNC_COMMIT)],
 )
 
-ENZYMEXLA_COMMIT = "1f193a884b1a6e149d43ed77a77a5bdb15882ba9"
+ENZYMEXLA_COMMIT = "fd5517f2223adcf579c165a0387231ef4931f55b"
 ENZYMEXLA_SHA256 = ""
 
 http_archive(

deps/build_local.jl

@@ -1,9 +1,49 @@
 # Invoke with
-# `julia --project=deps deps/build_local.jl [dbg/opt] [auto/cpu/cuda]`
+# `julia --project=deps deps/build_local.jl [--debug] [--backend=auto/cpu/cuda]`
 
 # the pre-built ReactantExtra_jll might not be loadable on this platform
 Reactant_jll = Base.UUID("0192cb87-2b54-54ad-80e0-3be72ad8a3c0")
 
+using ArgParse
+
+s = ArgParseSettings()
+#! format: off
+@add_arg_table! s begin
+    "--debug"
+        help = "Build with debug mode (-c dbg)."
+        action = :store_true
+    "--backend"
+        help = "Build with the specified backend (auto, cpu, cuda)."
+        default = "auto"
+        arg_type = String
+    "--gcc_host_compiler_path"
+        help = "Path to the gcc host compiler."
+        default = "/usr/bin/gcc"
+        arg_type = String
+    "--cc"
+        default = "/home/wmoses/llvms/llvm16-r/clang+llvm-16.0.2-x86_64-linux-gnu-ubuntu-22.04/bin/clang"
+        arg_type = String
+    "--hermetic_python_version"
+        help = "Hermetic Python version."
+        default = "3.10"
+        arg_type = String
+    # For GCC < 13 we need to disable these flags
+    "--xnn_disable_avx512fp16"
+        help = "Disable AVX512 FP16 support in XNNPACK."
+        action = :store_true
+    "--xnn_disable_avxvnniint8"
+        help = "Disable AVX VNNI INT8 support in XNNPACK."
+        action = :store_true
+end
+#! format: on
+parsed_args = parse_args(ARGS, s)
+
+println("Parsed args:")
+for (k, v) in parsed_args
+    println("  $k = $v")
+end
+println()
+
 using Pkg, Scratch, Preferences, Libdl
 
 # 1. Get a scratch directory
@@ -41,27 +81,10 @@ run(
 # --@local_config_cuda//:cuda_compiler=nvcc
 # --crosstool_top="@local_config_cuda//crosstool:toolchain"
 
-build_kind = if length(ARGS) ≥ 1
-    kind = ARGS[1]
-    if kind ∉ ("dbg", "opt")
-        error("Invalid build kind $(kind). Valid options are 'dbg' and 'opt'")
-    end
-    kind
-else
-    "dbg"
-end
+build_kind = parsed_args["debug"] ? "dbg" : "opt"
 
-@info "Building JLL with -c $(build_kind)"
-
-build_backend = if length(ARGS) ≥ 2
-    backend = ARGS[2]
-    if backend ∉ ("auto", "cpu", "cuda")
-        error("Invalid build backend $(backend). Valid options are 'auto', 'cpu', and 'cuda'")
-    end
-    backend
-else
-    "auto"
-end
+build_backend = parsed_args["backend"]
+@assert build_backend in ("auto", "cpu", "cuda")
 
 if build_backend == "auto"
     build_backend = try
@@ -78,8 +101,6 @@ elseif build_backend == "cpu"
     ""
 end
 
-@info "Building JLL with backend $(build_backend)"
-
 bazel_cmd = if !isnothing(Sys.which("bazelisk"))
     "bazelisk"
 elseif !isnothing(Sys.which("bazel"))
@@ -90,36 +111,28 @@ end
 
 @info "Building JLL with $(bazel_cmd)"
 
-if isempty(arg)
-    run(
-        Cmd(
-            `$(bazel_cmd) build -c $(build_kind) --action_env=JULIA=$(Base.julia_cmd().exec[1])
-            --repo_env HERMETIC_PYTHON_VERSION="3.10"
-            --check_visibility=false --verbose_failures :libReactantExtra.so`;
-            dir=source_dir,
-        ),
-    )
-else
-    run(
-        Cmd(
-            `$(bazel_cmd) build $(arg) -c $(build_kind) --action_env=JULIA=$(Base.julia_cmd().exec[1])
-            --repo_env=GCC_HOST_COMPILER_PATH=/usr/bin/gcc
-            --repo_env=CC=/home/wmoses/llvms/llvm16-r/clang+llvm-16.0.2-x86_64-linux-gnu-ubuntu-22.04/bin/clang
-            --repo_env HERMETIC_PYTHON_VERSION="3.10"
-            --check_visibility=false --verbose_failures :libReactantExtra.so`;
-            dir=source_dir,
-        ),
-    )
+gcc_host_compiler_path = parsed_args["gcc_host_compiler_path"]
+cc = parsed_args["cc"]
+hermetic_python_version = parsed_args["hermetic_python_version"]
+
+build_cmd_list = [bazel_cmd, "build"]
+!isempty(arg) && push!(build_cmd_list, arg)
+append!(build_cmd_list, ["-c", "$(build_kind)"])
+push!(build_cmd_list, "--action_env=JULIA=$(Base.julia_cmd().exec[1])")
+if parsed_args["xnn_disable_avx512fp16"]
+    push!(build_cmd_list, "--define=xnn_enable_avx512fp16=false")
+end
+if parsed_args["xnn_disable_avxvnniint8"]
+    push!(build_cmd_list, "--define=xnn_enable_avxvnniint8=false")
 end
-# env=Dict("HOME"=>ENV["HOME"], "PATH"=>joinpath(source_dir, "..")*":"*ENV["PATH"])))
+push!(build_cmd_list, "--repo_env=HERMETIC_PYTHON_VERSION=$(hermetic_python_version)")
+push!(build_cmd_list, "--repo_env=GCC_HOST_COMPILER_PATH=$(gcc_host_compiler_path)")
+push!(build_cmd_list, "--repo_env=CC=$(cc)")
+push!(build_cmd_list, "--check_visibility=false")
+push!(build_cmd_list, "--verbose_failures")
+push!(build_cmd_list, ":libReactantExtra.so")
 
-run(Cmd(`rm -f libReactantExtra.dylib`; dir=joinpath(source_dir, "bazel-bin")))
-run(
-    Cmd(
-        `ln -s libReactantExtra.so libReactantExtra.dylib`;
-        dir=joinpath(source_dir, "bazel-bin"),
-    ),
-)
+run(Cmd(Cmd(build_cmd_list); dir=source_dir))
 
 # Discover built libraries
 built_libs = filter(readdir(joinpath(source_dir, "bazel-bin"))) do file
@@ -129,7 +142,7 @@ end
 lib_path = joinpath(source_dir, "bazel-bin", only(built_libs))
 isfile(lib_path) || error("Could not find library $lib_path in build directory")
 
-# Tell ReactReactantExtra_jllant_jll to load our library instead of the default artifact one
+# Tell ReactantExtra_jll to load our library instead of the default artifact one
 set_preferences!(
     joinpath(dirname(@__DIR__), "LocalPreferences.toml"),
     "Reactant_jll",

docs/src/api/api.md

@@ -33,3 +33,10 @@ Reactant.Profiler.with_profiler
 Reactant.Profiler.annotate
 Reactant.Profiler.@annotate
 ```
+
+## Devices
+
+```@docs
+Reactant.devices
+Reactant.addressable_devices
+```

src/Compiler.jl

@@ -614,7 +614,13 @@ end
     @compile [optimize = ...] [no_nan = <true/false>] [sync = <true/false>] f(args...)
 """
 macro compile(args...)
-    default_options = Dict{Symbol,Any}(:optimize => true, :sync => false, :no_nan => false)
+    default_options = Dict{Symbol,Any}(
+        :optimize => true,
+        :sync => false,
+        :no_nan => false,
+        :client => nothing,
+        :device => nothing,
+    )
     return esc(first(compile_call_expr(__module__, compile, default_options, args...)))
 end
 
@@ -624,7 +630,13 @@ end
 Run @compile f(args..) then immediately execute it
 """
 macro jit(args...)
-    default_options = Dict{Symbol,Any}(:optimize => true, :sync => false, :no_nan => false)
+    default_options = Dict{Symbol,Any}(
+        :optimize => true,
+        :sync => false,
+        :no_nan => false,
+        :client => nothing,
+        :device => nothing,
+    )
     compile_expr, (; compiled, args) = compile_call_expr(
         __module__, compile, default_options, args...
     )
@@ -955,7 +967,7 @@ function codegen_xla_call(exec, flatten_names, donated_args_mask, nresults)
     return concretized_res_names, xla_call_code
 end
 
-function compile_xla(f, args; client=nothing, optimize=true, no_nan=false)
+function compile_xla(f, args; client=nothing, optimize=true, no_nan=false, device=nothing)
     # register MLIR dialects
     ctx = MLIR.IR.Context(Reactant.registry[], false)
     context_gc_vector[ctx] = Vector{TracedRArray}(undef, 0)
@@ -969,32 +981,69 @@ function compile_xla(f, args; client=nothing, optimize=true, no_nan=false)
             mod, f, args; optimize, no_nan
         )
 
-        if isnothing(client)
+        # Resolve client and device
+        device_ordinal = -1
+        if device === nothing
             if length(linear_args) > 0
-                for (k, _) in Iterators.filter(((_, v),) -> v isa TracedRArray, seen_args)
-                    client = XLA.client(k.data)
+                devices_list = [
+                    XLA.device(k.data) for (k, v) in seen_args if v isa TracedRArray
+                ]
+                if !isempty(devices_list)
+                    @assert allequal(devices_list) "All arguments must be on the same device: $(devices_list)"
+                    device = first(devices_list)
                 end
             end
-            if isnothing(client)
+        end
+
+        if client === nothing
+            if device !== nothing
+                client = XLA.client(device)
+            else
                 client = XLA.default_backend[]
+                device = XLA.ClientGetDevice(client, XLA.default_device_idx[])
+                device_ordinal = XLA.default_device_idx[]
+            end
+        else
+            if device !== nothing
+                @assert client == XLA.client(device) "client ($(client)) and XLA.client(device) ($(XLA.client(device))) must be the same"
+            else
+                device = XLA.ClientGetDevice(client, XLA.default_device_idx[])
+                device_ordinal = XLA.default_device_idx[]
             end
         end
 
+        if device_ordinal < 0
+            device_ordinal = XLA.DeviceToClientDeviceOrdinal(device)
+        end
+
         # compile MLIR module to XLA executable
-        exec = XLA.Compile(client, mod)
-        return exec,
-        linear_args, linear_results, preserved_args, seen_args, concrete_result,
-        isclosure
+        exec = XLA.Compile(
+            client,
+            mod;
+            device_ordinal,
+            num_replicas=1,
+            num_partitions=1,
+            use_shardy_partitioner=false,
+        )
+        return (
+            exec,
+            linear_args,
+            linear_results,
+            preserved_args,
+            seen_args,
+            concrete_result,
+            isclosure,
+        )
     finally
         MLIR.IR.deactivate!(ctx)
     end
     Base.delete!(context_gc_vector, ctx)
     return results
 end
 
-function compile(f, args; client=nothing, optimize=true, sync=false, no_nan=false)
+function compile(f, args; sync=false, kwargs...)
     exec, linear_args, linear_results, preserved_args, seen_args, concrete_result, isclosure = compile_xla(
-        f, args; client, optimize, no_nan
+        f, args; kwargs...
     )
 
     preserved_args_idx = last.(preserved_args)