feat: plugin API for new device implementations (#1241)

* refactor: move TPUs under accelerators

* feat: register using cplugin api from julia

* refactor: cleanup

* feat: metal plugin register

* fix: have different names for client

* fix: conditions

* fix: path to dylib

* fix: ordering

* fix: update C API

* chore: bump jll

* fix: api

* fix: use more recent metal plugins

* test: metal plugin tests

* ci: ensure compatibility

* fix: disable registration of metal for now

* chore: run fmt

* chore: bump jll

---------

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

Author: avik-pal

.buildkite/pipeline.yml

@@ -44,6 +44,7 @@ steps:
           REACTANT_TEST_GROUP: "{{matrix.group}}"
           JULIA_DEBUG: "Reactant,Reactant_jll"
           CUDA_VISIBLE_DEVICES: 0
+          REACTANT_BACKEND_GROUP: "GPU"
         if: build.message !~ /\[skip tests\]/
         timeout_in_minutes: 120
 
@@ -138,31 +139,3 @@ steps:
   #     rocm: "*"
   #   if: build.message !~ /\[skip tests\]/
   #   timeout_in_minutes: 60
-
-  # - label: "Metal Julia v{{matrix.version}}"
-  #   matrix:
-  #     setup:
-  #       version:
-  #         - "1.8"
-  #         - "1.9"
-  #   plugins:
-  #     - JuliaCI/julia#v1:
-  #         version: "{{matrix.version}}"
-  #   agents:
-  #     queue: "juliaecosystem"
-  #     os: "macos"
-  #     arch: "aarch64"
-  #   if: build.message !~ /\[skip tests\]/
-  #   timeout_in_minutes: 60
-  #   commands: |
-  #     echo "--- Setup Julia packages"
-  #     julia --color=yes -e '
-  #         import Pkg
-  #         Pkg.develop(; path = pwd())
-  #         Pkg.develop(; path = joinpath(pwd(), "lib", "EnzymeCore"))
-  #         Pkg.develop(; name = "Metal")' || exit 3
-
-  #     echo "+++ Run tests"
-  #     julia --color=yes test/metal.jl
-  #   env:
-  #     JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager

.github/workflows/CI.yml

@@ -143,6 +143,7 @@ jobs:
         id: run_tests
         env:
           JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
+          ENABLE_PJRT_COMPATIBILITY: 1
           REACTANT_TEST_GROUP: ${{ matrix.test_group }}
           XLA_FLAGS: "--xla_force_host_platform_device_count=12"
           JULIA_DEBUG: "Reactant,Reactant_jll"

deps/ReactantExtra/API.cpp

@@ -534,16 +534,8 @@ extern "C" PjRtClient *MakeTPUClient(const char *tpu_path, const char **error) {
     return nullptr;
   }
 
-  const PJRT_Api *pluginLoad =
-      LoadPjrtPlugin("tpu", tpu_library_path.c_str(), error);
-  if (pluginLoad == nullptr)
-    return nullptr;
-  auto tpu_status = InitializePjrtPlugin("tpu", error);
-  if (tpu_status)
-    return nullptr;
-
-  pjrt_client_register_profiler(pluginLoad);
-  return GetCApiClient("TPU");
+  return MakeClientUsingPluginAPI("tpu", tpu_library_path.c_str(), "TPU",
+                                  error);
 }
 
 extern "C" int ClientNumDevices(PjRtClient *client) {

src/Distributed.jl

@@ -1,6 +1,6 @@
 module Distributed
 
-using ..Reactant: Reactant
+using ..Reactant: Reactant, Accelerators
 using Sockets
 
 const initialized = Ref(false)
@@ -303,8 +303,8 @@ const _TPU_COORDINATOR_PORT = "8476"
 function get_coordinator_address(
     env::AbstractCloudTPUEnvDetector, timeout_in_seconds::Integer
 )
-    coordinator_address = if Reactant.TPUUtils.has_megascale_address()
-        Reactant.TPUUtils.get_tpu_env_value("MEGASCALE_COORDINATOR_ADDRESS")
+    coordinator_address = if Accelerators.TPU.has_megascale_address()
+        Accelerators.TPU.get_tpu_env_value("MEGASCALE_COORDINATOR_ADDRESS")
     else
         first(_get_worker_list_in_slice(env))
     end
@@ -361,13 +361,13 @@ function get_process_id(env::AbstractCloudTPUEnvDetector)
 end
 
 function _get_num_slices(::AbstractCloudTPUEnvDetector)
-    Reactant.TPUUtils.has_megascale_address() || return 1
-    return parse(Int, Reactant.TPUUtils.get_tpu_env_value("MEGASCALE_NUM_SLICES"))
+    Accelerators.TPU.has_megascale_address() || return 1
+    return parse(Int, Accelerators.TPU.get_tpu_env_value("MEGASCALE_NUM_SLICES"))
 end
 
 function _get_slice_id(::AbstractCloudTPUEnvDetector)
-    Reactant.TPUUtils.has_megascale_address() || return 0
-    return parse(Int, Reactant.TPUUtils.get_tpu_env_value("MEGASCALE_SLICE_ID"))
+    Accelerators.TPU.has_megascale_address() || return 0
+    return parse(Int, Accelerators.TPU.get_tpu_env_value("MEGASCALE_SLICE_ID"))
 end
 
 function _get_process_id_in_slice end
@@ -376,7 +376,7 @@ function _get_worker_list_in_slice end
 ## GceTPUCluster
 
 function is_env_present(::GceTPUCluster)
-    if !Reactant.TPUUtils.RUNNING_IN_CLOUD_TPU_VM[]
+    if !Accelerators.TPU.RUNNING_IN_CLOUD_TPU_VM[]
         @debug "Did not detect cloud TPU VM"
         return false
     end
@@ -386,8 +386,8 @@ function is_env_present(::GceTPUCluster)
         return false
     end
 
-    metadata_response, metadata_code = Reactant.TPUUtils.get_metadata("agent-worker-number")
-    if metadata_code == Reactant.TPUUtils._TPU_METADATA_RESPONSE_CODE_SUCCESS
+    metadata_response, metadata_code = Accelerators.TPU.get_metadata("agent-worker-number")
+    if metadata_code == Accelerators.TPU._TPU_METADATA_RESPONSE_CODE_SUCCESS
         @debug "Gce Tpu Cluster detected for Reactant Distributed System"
         return true
     else
@@ -400,23 +400,23 @@ function is_env_present(::GceTPUCluster)
 end
 
 function _get_process_id_in_slice(::GceTPUCluster)
-    return parse(Int, first(Reactant.TPUUtils.get_metadata("agent-worker-number")))
+    return parse(Int, first(Accelerators.TPU.get_metadata("agent-worker-number")))
 end
 
 function _get_worker_list_in_slice(::GceTPUCluster)
-    workers = split(first(Reactant.TPUUtils.get_metadata("worker-network-endpoints")), ',')
+    workers = split(first(Accelerators.TPU.get_metadata("worker-network-endpoints")), ',')
     return [split(w, ':')[3] for w in workers]
 end
 
 ## GkeTPUCluster
 
 function is_env_present(::GkeTPUCluster)
-    if Reactant.TPUUtils.RUNNING_IN_CLOUD_TPU_VM[] && haskey(ENV, "TPU_WORKER_HOSTNAMES")
+    if Accelerators.TPU.RUNNING_IN_CLOUD_TPU_VM[] && haskey(ENV, "TPU_WORKER_HOSTNAMES")
         @debug "Detected GKE TPU cluster for Reactant Distributed System"
         return true
     end
 
-    if !Reactant.TPUUtils.RUNNING_IN_CLOUD_TPU_VM[]
+    if !Accelerators.TPU.RUNNING_IN_CLOUD_TPU_VM[]
         @debug "Did not detect cloud TPU VM"
         return false
     end

src/Reactant.jl

@@ -45,9 +45,9 @@ function ancestor(T::Type{<:AbstractArray})
     return T
 end
 
-include("TPUs.jl")
+include("accelerators/Accelerators.jl")
 
-using .TPUUtils: has_tpu
+using .Accelerators.TPU: has_tpu
 
 include("mlir/MLIR.jl")
 include("xla/XLA.jl")

src/accelerators/Accelerators.jl

@@ -0,0 +1,6 @@
+module Accelerators
+
+include("TPU.jl")
+include("Metal.jl")
+
+end

src/accelerators/Metal.jl

@@ -0,0 +1,60 @@
+module Metal
+
+using Reactant: Reactant
+using Scratch: @get_scratch!
+using Downloads
+
+const metal_pjrt_plugin_dir = Ref{Union{Nothing,String}}(nothing)
+
+function __init__()
+    @static if Sys.isapple()
+        Reactant.precompiling() || setup_metal_pjrt_plugin!()
+    end
+end
+
+function setup_metal_pjrt_plugin!()
+    path_from_env = get(ENV, "METAL_LIBRARY_PATH", nothing)
+    if path_from_env !== nothing && ispath(path_from_env)
+        metal_pjrt_plugin_dir[] = path_from_env
+    else
+        metal_pjrt_plugin_dir[] = @get_scratch!("pjrt_metal_plugin")
+    end
+    download_metal_pjrt_plugin_if_needed(metal_pjrt_plugin_dir[])
+    return nothing
+end
+
+get_metal_pjrt_plugin_dir() = metal_pjrt_plugin_dir[]
+
+function get_metal_pjrt_plugin_path()
+    return joinpath(get_metal_pjrt_plugin_dir(), "pjrt_plugin_metal_14.dylib")
+end
+
+function download_metal_pjrt_plugin_if_needed(path=nothing)
+    path === nothing && (path = get_metal_pjrt_plugin_dir())
+    @assert path !== nothing "metal_pjrt_plugin_dir is not set!"
+
+    metal_pjrt_plugin_path = joinpath(path, "pjrt_plugin_metal_14.dylib")
+    if !isfile(metal_pjrt_plugin_path)
+        zip_file_path = joinpath(path, "pjrt-plugin-metal.zip")
+        tmp_dir = joinpath(path, "tmp")
+        Downloads.download(
+            if Sys.ARCH === :aarch64
+                "https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_metal-0.1.1-py3-none-macosx_13_0_arm64.whl"
+            elseif Sys.ARCH === :x86_64
+                "https://files.pythonhosted.org/packages/87/ec/9bb7f7f0ffd06c3fb89813126b2f698636ac7a4263ed7bdd1ff7d7c94f8f/jax_metal-0.1.1-py3-none-macosx_10_14_x86_64.whl"
+            else
+                error("Unsupported architecture: $(Sys.ARCH)")
+            end,
+            zip_file_path,
+        )
+        run(`unzip -qq $(zip_file_path) -d $(tmp_dir)`)
+        mv(
+            joinpath(tmp_dir, "jax_plugins", "metal_plugin", "pjrt_plugin_metal_14.dylib"),
+            metal_pjrt_plugin_path,
+        )
+        rm(tmp_dir; recursive=true)
+        rm(zip_file_path; recursive=true)
+    end
+end
+
+end

src/TPUs.jl renamed to src/accelerators/TPU.jl

@@ -1,4 +1,4 @@
-module TPUUtils
+module TPU
 
 using Reactant: Reactant
 using EnumX: @enumx
@@ -31,20 +31,24 @@ end
 
 get_libtpu_dir() = libtpu_dir[]
 
-get_libtpu_path() = get_libtpu_dir() * "/libtpu.so"
+get_libtpu_path() = joinpath(get_libtpu_dir(), "libtpu.so")
 
 function download_libtpu_if_needed(path=nothing)
     path === nothing && (path = get_libtpu_dir())
     @assert path !== nothing "libtpu_dir is not set!"
-    if !isfile(path * "/libtpu.so")
+
+    libtpu_path = joinpath(path, "libtpu.so")
+    if !isfile(libtpu_path)
+        zip_file_path = joinpath(path, "tpu.zip")
+        tmp_dir = joinpath(path, "tmp")
         Downloads.download(
             "https://storage.googleapis.com/libtpu-nightly-releases/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20250415+nightly-py3-none-manylinux_2_31_x86_64.whl",
-            path * "/tpu.zip",
+            zip_file_path,
         )
-        run(`unzip -qq $(path*"/tpu.zip") -d $(path)/tmp`)
-        run(`mv $(path)/tmp/libtpu/libtpu.so $(path)/libtpu.so`)
-        rm(path * "/tmp"; recursive=true)
-        rm(path * "/tpu.zip"; recursive=true)
+        run(`unzip -qq $(zip_file_path) -d $(tmp_dir)`)
+        mv(joinpath(tmp_dir, "libtpu", "libtpu.so"), libtpu_path)
+        rm(tmp_dir; recursive=true)
+        rm(zip_file_path; recursive=true)
     end
 end
 

src/xla/IFRT/Client.jl

@@ -112,13 +112,15 @@ end
 
 # Different Backends
 const cpu_client_count = Ref(0)
-const gpu_client_count = Ref(0)
+const cuda_client_count = Ref(0)
 const tpu_client_count = Ref(0)
+const metal_client_count = Ref(0)
 
 for (backend, counter) in (
     (:CPUClient, :cpu_client_count),
-    (:GPUClient, :gpu_client_count),
+    (:CUDAClient, :cuda_client_count),
     (:TPUClient, :tpu_client_count),
+    (:MetalClient, :metal_client_count),
 )
     main_fn = Symbol(:MakeIFRTPJRT, backend)
     @eval function $(backend)(args...; checkcount::Bool=true, kwargs...)
@@ -159,7 +161,7 @@ function MakeIFRTPJRTCPUClient(;
     return client, refstr
 end
 
-function MakeIFRTPJRTGPUClient(;
+function MakeIFRTPJRTCUDAClient(;
     node_id::Integer=0,
     num_nodes::Integer=1,
     platform::String="gpu",
@@ -196,19 +198,51 @@ function MakeIFRTPJRTTPUClient(;
     num_nodes::Integer=1,
     distributed_runtime_client::Union{Nothing,XLA.DistributedRuntimeClient}=nothing,
 )
-    refstr = Ref{Cstring}()
+    return MakeIFRTPJRTClientViaPluginAPI(
+        tpu_path, "tpu", "TPU"; node_id, num_nodes, distributed_runtime_client
+    )
+end
+
+function MakeIFRTPJRTMetalClient(;
+    metal_pjrt_plugin_path::String,
+    node_id::Integer=0,
+    num_nodes::Integer=1,
+    distributed_runtime_client::Union{Nothing,XLA.DistributedRuntimeClient}=nothing,
+)
+    return MakeIFRTPJRTClientViaPluginAPI(
+        metal_pjrt_plugin_path,
+        "metal",
+        "METAL";
+        node_id,
+        num_nodes,
+        distributed_runtime_client,
+    )
+end
+
+function MakeIFRTPJRTClientViaPluginAPI(
+    library_path::String,
+    device_type::String,
+    client_name::String=uppercase(device_type);
+    node_id::Integer=0,
+    num_nodes::Integer=1,
+    distributed_runtime_client::Union{Nothing,XLA.DistributedRuntimeClient}=nothing,
+)
+    pjrt_client = XLA.PJRT.MakeClientUsingPluginAPI(library_path, device_type, client_name)
+
     distributed_runtime_client =
         distributed_runtime_client === nothing ? C_NULL : distributed_runtime_client.client
 
-    GC.@preserve refstr distributed_runtime_client begin
-        client = @ccall MLIR.API.mlir_c.ifrt_make_pjrt_tpu_client(
-            tpu_path::Cstring,
-            refstr::Ptr{Cstring},
+    errstr = Ref{Cstring}()
+    GC.@preserve pjrt_client errstr distributed_runtime_client device_type begin
+        client = @ccall MLIR.API.mlir_c.ifrt_pjrt_make_client_with_default_kv_store(
+            pjrt_client::Ptr{Cvoid},
             node_id::Cint,
             num_nodes::Cint,
             distributed_runtime_client::Ptr{Cvoid},
+            errstr::Ptr{Cstring},
+            device_type::Cstring,
         )::Ptr{Cvoid}
     end
 
-    return client, refstr
+    return client, errstr
 end