Create TT client via plugin for TensTorrent devices

Author: giordano

src/accelerators/Accelerators.jl

@@ -2,5 +2,6 @@ module Accelerators
 
 include("TPU.jl")
 include("Metal.jl")
+include("TT.jl")
 
 end

src/accelerators/TT.jl

@@ -0,0 +1,62 @@
+module TT
+
+using Reactant: Reactant
+using Scratch: @get_scratch!
+using Downloads
+
+const tt_pjrt_plugin_dir = Ref{Union{Nothing,String}}(nothing)
+
+function __init__()
+    @static if Sys.islinux()
+        Reactant.precompiling() || setup_tt_pjrt_plugin!()
+    end
+end
+
+has_tt() = true
+
+function setup_tt_pjrt_plugin!()
+    path_from_env = get(ENV, "TT_LIBRARY_PATH", nothing)
+    if path_from_env !== nothing && ispath(path_from_env)
+        tt_pjrt_plugin_dir[] = path_from_env
+    else
+        tt_pjrt_plugin_dir[] = @get_scratch!("pjrt_tt_plugin")
+    end
+    # download_tt_pjrt_plugin_if_needed(tt_pjrt_plugin_dir[])
+    return nothing
+end
+
+get_tt_pjrt_plugin_dir() = tt_pjrt_plugin_dir[]
+
+function get_tt_pjrt_plugin_path()
+    return joinpath(get_tt_pjrt_plugin_dir(), "pjrt_plugin_tt.so")
+end
+
+# function download_tt_pjrt_plugin_if_needed(path=nothing)
+#     path === nothing && (path = get_tt_pjrt_plugin_dir())
+#     @assert path !== nothing "tt_pjrt_plugin_dir is not set!"
+
+#     tt_pjrt_plugin_path = joinpath(path, "pjrt_plugin_tt_14.dylib")
+#     if !isfile(tt_pjrt_plugin_path)
+#         zip_file_path = joinpath(path, "pjrt-plugin-tt.zip")
+#         tmp_dir = joinpath(path, "tmp")
+#         Downloads.download(
+#             if Sys.ARCH === :aarch64
+#                 "https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_tt-0.1.1-py3-none-macosx_13_0_arm64.whl"
+#             elseif Sys.ARCH === :x86_64
+#                 "https://files.pythonhosted.org/packages/87/ec/9bb7f7f0ffd06c3fb89813126b2f698636ac7a4263ed7bdd1ff7d7c94f8f/jax_tt-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", "tt_plugin", "pjrt_plugin_tt_14.dylib"),
+#             tt_pjrt_plugin_path,
+#         )
+#         rm(tmp_dir; recursive=true)
+#         rm(zip_file_path; recursive=true)
+#     end
+# end
+
+end # module TT

src/xla/IFRT/Client.jl

@@ -115,12 +115,14 @@ const cpu_client_count = Ref(0)
 const cuda_client_count = Ref(0)
 const tpu_client_count = Ref(0)
 const metal_client_count = Ref(0)
+const tt_client_count = Ref(0)
 
 for (backend, counter) in (
     (:CPUClient, :cpu_client_count),
     (:CUDAClient, :cuda_client_count),
     (:TPUClient, :tpu_client_count),
     (:MetalClient, :metal_client_count),
+    (:TTClient, :tt_client_count),
 )
     main_fn = Symbol(:MakeIFRTPJRT, backend)
     @eval function $(backend)(args...; checkcount::Bool=true, kwargs...)
@@ -219,6 +221,22 @@ function MakeIFRTPJRTMetalClient(;
     )
 end
 
+function MakeIFRTPJRTTTClient(;
+    rocm_pjrt_plugin_path::String,
+    node_id::Integer=0,
+    num_nodes::Integer=1,
+    distributed_runtime_client::Union{Nothing,XLA.DistributedRuntimeClient}=nothing,
+)
+    return MakeIFRTPJRTClientViaPluginAPI(
+        rocm_pjrt_plugin_path,
+        "rocm",
+        "TT";
+        node_id,
+        num_nodes,
+        distributed_runtime_client,
+    )
+end
+
 function MakeIFRTPJRTClientViaPluginAPI(
     library_path::String,
     device_type::String,

src/xla/PJRT/Client.jl

@@ -110,12 +110,14 @@ const cpu_client_count = Ref(0)
 const cuda_client_count = Ref(0)
 const tpu_client_count = Ref(0)
 const metal_client_count = Ref(0)
+const tt_client_count = Ref(0)
 
 for (backend, counter) in (
     (:CPUClient, :cpu_client_count),
     (:CUDAClient, :cuda_client_count),
     (:TPUClient, :tpu_client_count),
     (:MetalClient, :metal_client_count),
+    (:TTClient, :tt_client_count),
 )
     main_fn = Symbol(:Make, backend)
     @eval function $(backend)(args...; checkcount::Bool=true, kwargs...)
@@ -207,6 +209,20 @@ function MakeMetalClient(;
     return MakeClientUsingPluginAPI(metal_pjrt_plugin_path, "metal", "METAL")
 end
 
+function MakeTTClient(;
+    tt_pjrt_plugin_path::String,
+    node_id::Integer=0,
+    num_nodes::Integer=1,
+    distributed_runtime_client::Union{Nothing,XLA.DistributedRuntimeClient}=nothing,
+)
+    @assert node_id == 0 "`PJRT.MakeTTClient` does not support node_id"
+    @assert num_nodes == 1 "`PJRT.MakeTTClient` does not support num_nodes > 1"
+    @assert distributed_runtime_client === nothing "`PJRT.MakeTTClient` does not support \
+                                                    distributed_runtime_client"
+
+    return MakeClientUsingPluginAPI(tt_pjrt_plugin_path, "tt", "TT")
+end
+
 function MakeClientUsingPluginAPI(
     library_path::String, device_type::String, client_name::String=uppercase(device_type)
 )

src/xla/XLA.jl

@@ -226,6 +226,32 @@ for runtime in (:PJRT, :IFRT)
                     catch e
                         println(stdout, e)
                     end
+                elseif Accelerators.TT.has_tt()
+                    try
+                        if was_initialized && haskey(state.clients, "tt")
+                            XLA.free_client(state.clients["tt"])
+                            XLA.$(runtime).tt_client_count[] -= 1
+                        end
+                        # The env var `TT_METAL_RUNTIME_ROOT` must be set before creating the client.
+                        if isnothing(get(ENV, "TT_METAL_RUNTIME_ROOT", nothing))
+                            tt_metal_path_in_wheel = joinpath(dirname(Accelerators.TT.get_tt_pjrt_plugin_path()), "tt-metal")
+                            if ispath(tt_metal_path_in_wheel)
+                                ENV["TT_METAL_RUNTIME_ROOT"] = tt_metal_path_in_wheel
+                            else
+                                error("`TT_METAL_RUNTIME_ROOT` environment variable not set and we could not automatically determine it")
+                            end
+                        end
+
+                        tt = $(runtime).TTClient(
+                            ;
+                            tt_pjrt_plugin_path=Accelerators.TT.get_tt_pjrt_plugin_path(),
+                            common_kwargs...
+                                )
+                        state.clients["tt"] = tt
+                        state.default_client = tt
+                    catch e
+                            println(stdout, e)
+                    end
                 elseif Reactant_jll.host_platform.tags["gpu"] != "none"
                     try
                         if was_initialized && haskey(state.clients, "cuda")