Does EXLA support the Apple Silicon GPU (MPS)?

Elixir Framework Forums Nx Forum
,
1

Hi!,

I have the Mac Studio and I want to use the GPU on it with EXLA.
I made the project that has config like below.

import Config
config :nx, :default_backend, EXLA.Backend

with iex ,
Nx.tensor [1,2,3,4] ,
it show
[info] TfrtCpuClient created.
nx.Tensor<
s64[4]
EXLA.Backend<host:0, 0.1189406505.872022046.161358>
[1, 2, 3, 4]

it seems not use GPU.

And, with below code (to make long Nx calculation) ----
url =
“https://pix4free.org/assets/library/2021-01-12/originals/san_francisco_california_golden_gate_bridge_water.jpg”

%{body: raw_image} = Req.get!(url)

imageNx = StbImage.read_binary!(raw_image)
            |> StbImage.to_nx()
            |> Nx.reshape({:auto,3})

IO.puts "Started --------"
start_ts =DateTime.utc_now() |> DateTime.to_unix(:second)

KMeans.fit(imageNx,
    num_clusters: 10,
    num_runs: 10,
    max_iterations: 200,
    key: Nx.Random.key(42)
  )

end_ts =DateTime.utc_now() |> DateTime.to_unix(:second)
IO.puts "DONE -- ellapse : #{end_ts - start_ts}"

It takes about 230 seconds, and the Mac Activity Monitor show 0 GPU usage of beam_smp.

Does the EXLA support Apple Silicon GPU ?

Any suggestion will be appreciated…

Thanks,

1 Like
2

We don’t currently support mac GPUs unfortunately.
There’s ongoing work on migrating EXLA to an MLIR-based implementation that will in the future enable support, though.

Torchx is closer to that, altough Pytorch still has a few issues to go on the ways of API coverage for Apple Metal.

2 Likes
3

Ah, there’s something else you should also configure, the default defn compiler!

config :nx, :default_defn_options, compiler: EXLA should make your code quite a bit faster too

4 Likes
4

Thank you so much!

Following your advice, adding the “config :nx, :default_defn_options, compiler: EXLA” setting dramatically reduced the computation time.

It went down from 220 seconds to 13 seconds. That’s really amazing.

Thank you once again.

4 Likes
5

Awesome!