Skip to content

fix: Default to the CPU device when Metal is not available - #148

Open
georgeguimaraes wants to merge 2 commits into
elixir-nx:mainfrom
georgeguimaraes:fix/compiler-default-device
Open

georgeguimaraes wants to merge 2 commits into
elixir-nx:mainfrom
georgeguimaraes:fix/compiler-default-device

Conversation

@georgeguimaraes

Copy link
Copy Markdown
Member

EMLX.default_device/0 returns :gpu when :default_device is not configured. On Linux, EMLX.Application does not start a GPU worker, so a jit with compiler: EMLX and no device: option fails. Users have to pass device: :cpu to every compile, or set the config, even though the CPU is the only device.

I found this in CI on Linux, where the same code worked on a Mac.

With this change, default_device/0 returns :cpu when the application did not start a GPU worker. A configured :default_device still wins, and on a Mac with Metal the default is still :gpu.

EMLX.default_device/0 returned :gpu when :default_device was not configured. On Linux there is no GPU worker, so a jit with compiler: EMLX and no device: option failed.

Now default_device/0 returns :cpu when the application did not start a GPU worker. A configured :default_device still wins.
Comment thread emlx/lib/emlx.ex
Comment on lines 2545 to 2554
def default_device do
Application.get_env(:emlx, :default_device, :gpu)
case Application.get_env(:emlx, :default_device) do
nil ->
# EMLX.Application starts a GPU worker only when Metal is available.
if :persistent_term.get({EMLX, :default_worker, :gpu}, nil), do: :gpu, else: :cpu

device ->
device
end
end

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is not the correct solution and should be solved at the runtime config level.
As this stands, there can be a race condition where an app runs EMLX.default_device before the default worker is started and therefore gets cpu instead of gpu.

We should expose mlx::core::is_available as a NIF and call it with the result of Application.get_env(:emlx, :default_device, :gpu). If the configured default device is not available, we return :cpu instead

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants