|
1 | 1 | # EMLX |
2 | 2 |
|
3 | | -[](https://hex.pm/packages/emlx) [](https://hexdocs.pm/emlx) |
| 3 | +This repository holds the following projects: |
4 | 4 |
|
5 | | -EMLX is the Nx Backend for the [MLX](https://github.com/ml-explore/mlx) library. |
| 5 | +* [`EMLX`](https://github.com/elixir-nx/emlx/tree/main/emlx#readme) - Nx backend for Apple's [MLX](https://github.com/ml-explore/mlx) library, enabling GPU-accelerated tensor operations on Apple Silicon |
6 | 6 |
|
7 | | -Because of MLX's nature, EMLX with GPU backend is only supported on macOS. |
| 7 | +* [`EMLXAxon`](https://github.com/elixir-nx/emlx/tree/main/emlx_axon#readme) - [Axon](https://github.com/elixir-nx/axon) model rewrites that swap supported nodes for `EMLX.Fast` Metal shader implementations, accelerating LLM inference on Apple Silicon |
8 | 8 |
|
9 | | -MLX with CPU backend is available on most mainstream platforms, however, the CPU backend may not be as optimized as the GPU backend, |
10 | | -especially for non-macOS OSes, as they're not prioritized for development. Right now, EMLX supports x86_64 and arm64 architectures |
11 | | -on both macOS and Linux. |
| 9 | +Each has their own README, which you can access above to learn more. |
12 | 10 |
|
13 | | -The M-Series Macs have an unified memory architecture, which allows for more passing data between the CPU and GPU to be effectively a no-op. |
14 | | - |
15 | | -Besides the backend, this library also provides a `Nx.Defn.Compiler` implementation that JIT-compiles Nx functions with smart use of MLX command queues. |
16 | | - |
17 | | -- **Worker-thread dispatch** — MLX ops run on dedicated threads instead of BEAM dirty schedulers, eliminating scheduler starvation under load. |
18 | | -- **Per-process Metal command queues (`EMLX.CommandQueue`)** — each BEAM process can get its own GPU command queue for true process-level GPU isolation. |
19 | | -- **GPU pointer interop** — `Nx.Backend.from_pointer/5` and `to_pointer/2` for zero-copy Metal buffer sharing with other languages, such as with Python via Pythonx. |
20 | | - |
21 | | -Metal does not support 64-bit floats, so neither MLX nor EMLX do either. |
22 | | - |
23 | | -## Usage |
24 | | - |
25 | | -To use EMLX, you can add it as a dependency in your `mix.exs`: |
26 | | - |
27 | | -```elixir |
28 | | -def deps do |
29 | | - [ |
30 | | - {:emlx, github: "elixir-nx/emlx", branch: "main"} |
31 | | - ] |
32 | | -end |
33 | | -``` |
34 | | - |
35 | | -Then, you just need to set `EMLX.Backend` as the default backend for your Nx functions: |
36 | | - |
37 | | -```elixir |
38 | | -Nx.default_backend(EMLX.Backend) |
39 | | - |
40 | | -# Setting the device to the CPU (default) |
41 | | -Nx.default_backend({EMLX.Backend, device: :cpu}) |
42 | | - |
43 | | -# Setting the device to the GPU |
44 | | -Nx.default_backend({EMLX.Backend, device: :gpu}) |
45 | | - |
46 | | -# or use the application config using one of the alternatives above as the value: |
47 | | - |
48 | | -config :nx, :default_backend, EMLX.Backend |
49 | | -config :nx, :default_backend, {EMLX.Backend, device: :cpu} |
50 | | -config :nx, :default_backend, {EMLX.Backend, device: :gpu} |
51 | | -``` |
52 | | - |
53 | | -If you want to use the JIT compiler, you can set the default compiler as shown below. |
54 | | - |
55 | | -```elixir |
56 | | -Nx.Defn.default_options(compiler: EMLX) |
57 | | - |
58 | | -# Alternatively, we can set this in the application environment |
59 | | - |
60 | | -config :nx, :default_defn_options, compiler: EMLX |
61 | | -``` |
62 | | - |
63 | | -### MLX binaries |
64 | | - |
65 | | -EMLX relies on the [MLX](https://github.com/ml-explore/mlx) library to function, and currently EMLX will download precompiled builds from [mlx-build](https://github.com/cocoa-xu/mlx-build). |
66 | | - |
67 | | -#### Using precompiled binaries |
68 | | - |
69 | | -While the default configuration should be suitable for most cases, there is however a number of environment variables that you may want to use in order to customize the variant of MLX binary. |
70 | | - |
71 | | -The binaries are always downloaded to match the current configuration, so you should set the environment variables in .bash_profile or a similar configuration file so you don't need to export it in every shell session. |
72 | | - |
73 | | -##### `LIBMLX_VERSION` |
74 | | - |
75 | | -The version of the MLX binary to download. By default EMLX will always use the latest version possible. |
76 | | - |
77 | | -##### `LIBMLX_ENABLE_JIT` |
78 | | - |
79 | | -Defaults to `false`. |
80 | | - |
81 | | -Using JIT compilation for Metal kernels when set to `true`. |
82 | | - |
83 | | -##### `LIBMLX_ENABLE_DEBUG` |
84 | | - |
85 | | -Defaults to `false`. |
86 | | - |
87 | | -Enhance metal debug workflow by enabling debug information in the Metal shaders when set to `true`. |
88 | | - |
89 | | -##### `LIBMLX_CACHE` |
90 | | - |
91 | | -The directory to store the downloaded and built archives in. Defaults to the standard cache location for the given operating system. |
92 | | - |
93 | | -#### Compiling from source |
94 | | - |
95 | | -If you want to compile MLX from source, you can do so by setting the `LIBMLX_BUILD` environment variable to `true`. |
96 | | - |
97 | | -Environment variables listed in the previous section will still apply. |
| 11 | +[Check our organization page for a general introduction to Machine Learning in Elixir](https://github.com/elixir-nx/). |
0 commit comments