mlx Documentation

MLX: An array framework for Apple silicon

27,926 stars C++ C++ 1 file · ~14,152 tokens #mlx
RAW Doc

File: README.md

MLX

Quickstart | Installation |
Documentation |
Examples

MLX is an array framework for machine learning on Apple silicon,
brought to you by Apple machine learning research.

Some key features of MLX include:

  • Familiar APIs: MLX has a Python API that closely follows NumPy. MLX
    also has fully featured C++, C, and
    Swift APIs, which closely mirror
    the Python API. MLX has higher-level packages like mlx.nn and
    mlx.optimizers with APIs that closely follow PyTorch to simplify building
    more complex models.

  • Composable function transformations: MLX supports composable function
    transformations for automatic differentiation, automatic vectorization,
    and computation graph optimization.

  • Lazy computation: Computations in MLX are lazy. Arrays are only
    materialized when needed.

  • Dynamic graph construction: Computation graphs in MLX are constructed
    dynamically. Changing the shapes of function arguments does not trigger
    slow compilations, and debugging is simple and intuitive.

  • Multi-device: Operations can run on any of the supported devices
    (currently the CPU and the GPU).

  • Unified memory: A notable difference from MLX and other frameworks
    is the unified memory model. Arrays in MLX live in shared memory.
    Operations on MLX arrays can be performed on any of the supported
    device types without transferring data.

MLX is designed by machine learning researchers for machine learning
researchers. The framework is intended to be user-friendly, but still efficient
to train and deploy models. The design of the framework itself is also
conceptually simple. We intend to make it easy for researchers to extend and
improve MLX with the goal of quickly exploring new ideas.

The design of MLX is inspired by frameworks like
NumPy,
PyTorch, Jax, and
ArrayFire.

Examples

The MLX examples repo has a
variety of examples, including:

Quickstart

See the quick start
guide
in the documentation.

Installation

MLX is available on PyPI. To install MLX on
macOS, run:

bash
pip install mlx

To install the CUDA backend on Linux, run:

bash
pip install mlx[cuda]

To install a CPU-only Linux package, run:

bash
pip install mlx[cpu]

Checkout the
documentation
for more information on building the C++ and Python APIs from source.

Contributing

Check out the contribution guidelines for more information
on contributing to MLX. See the
docs for more
information on building from source, and running tests.

We are grateful for all of our
contributors. If you contribute
to MLX and wish to be acknowledged, please add your name to the list in your
pull request.

Citing MLX

The MLX software suite was initially developed with equal contribution by Awni
Hannun, Jagrit Digani, Angelos Katharopoulos, and Ronan Collobert. If you find
MLX useful in your research and wish to cite it, please use the following
BibTex entry:

text
@software{mlx2023,
  author = {Awni Hannun and Jagrit Digani and Angelos Katharopoulos and Ronan Collobert},
  title = {{MLX}: Efficient and flexible machine learning on Apple silicon},
  url = {https://github.com/ml-explore},
  version = {0.0},
  year = {2023},
}

File: benchmarks/python/comparative/README.md

Microbenchmarks comparing MLX to PyTorch

Implement the same microbenchmarks in MLX and PyTorch to compare and make a
list of the biggest possible performance improvements and/or regressions.

Run with python bench_mlx.py sum_axis --size 8x1024x128 --axis 2 --cpu for
instance to measure the times it takes to sum across the 3rd axis of the above
tensor on the cpu.

compare.py runs several benchmarks and compares the speed-up or lack thereof
in comparison to PyTorch.

Each bench script can be run with --print-pid to print the PID and wait for a
key in order to ease attaching a debugger.

File: docs/src/_templates/module-base-class.rst

{{ fullname | escape | underline}}

.. currentmodule:: {{ module }}

.. add toctree option to make autodoc generate the pages

.. autoclass:: {{ objname }}

{% block attributes %}
{% if attributes %}
.. rubric:: Attributes

.. autosummary::
:toctree: .
{% for item in attributes %}
~{{ fullname }}.{{ item }}
{%- endfor %}
{% endif %}
{% endblock %}

{% block methods %}
{% if methods %}
.. rubric:: Methods

.. autosummary::
:toctree: .
{% for item in methods %}
{%- if item not in inherited_members and item != 'init' %}
~{{ fullname }}.{{ item }}
{%- endif -%}
{%- endfor %}
{% endif %}
{% endblock %}

File: docs/src/_templates/nn-module-template.rst

{{ fullname | escape | underline}}

.. currentmodule:: {{ module }}

.. autoclass:: {{ objname }}

{% block methods %}

{% if methods %}
.. rubric:: {{ _('Methods') }}

.. autosummary::
{% for item in methods %}
{%- if item not in inherited_members and item != "init" %}
~{{ name }}.{{ item }}
{%- endif %}
{%- endfor %}
{% endif %}
{% endblock %}

File: docs/src/_templates/optimizers-template.rst

{{ fullname | escape | underline}}

.. currentmodule:: {{ module }}

.. autoclass:: {{ objname }}

{% block methods %}

{% if methods %}
.. rubric:: {{ _('Methods') }}

.. autosummary::
{% for item in methods %}
{%- if item not in inherited_members %}
~{{ name }}.{{ item }}
{%- endif %}
{%- endfor %}
{% endif %}
{% endblock %}

File: docs/src/dev/custom_metal_kernels.rst

.. _custom_metal_kernels:

Custom Metal Kernels

MLX supports writing custom Metal kernels through the Python and C++ APIs.

Simple Example

.. currentmodule:: mlx.core

Let's write a custom kernel that computes exp elementwise:

.. code-block:: python

source = """
uint elem = thread_position_in_grid.x;
T tmp = inp[elem];
out[elem] = metal::exp(tmp);
"""

kernel = mx.fast.metal_kernel(
name="myexp",
input_names=["inp"],
output_names=["out"],
source=source,
)

def exp_elementwise(a: mx.array):
outputs = kernel(
inputs=[a],
template=[("T", mx.float32)],
grid=(a.size, 1, 1),
threadgroup=(256, 1, 1),
output_shapes=[a.shape],
output_dtypes=[a.dtype],
)
return outputs[0]

a = mx.random.normal(shape=(4, 16)).astype(mx.float16)
b = exp_elementwise(a)
assert mx.allclose(b, mx.exp(a))

Every time you make a kernel, a new Metal library is created and possibly
JIT compiled. To reduce the overhead from that, build the kernel once with
:func:fast.metal_kernel and then use it many times.

.. note::
Only pass the body of the Metal kernel in source. The function
signature is generated automatically.

The full function signature will be generated using:

  • The shapes/dtypes of inputs
    In the above, a is an mx.array of type mx.float16 and we pass it with the key inp
    so we will add const device float16_t* inp to the signature.
    inp_shape, inp_strides and inp_ndim are also added for convenience if they are present
    in source.
  • The list of output_dtypes
    In the above, out is an mx.array of type mx.float16
    so we add device float16_t* out.
  • Template parameters passed using template
    In the above, template=[("T", mx.float32)] adds a template of template <typename T> to the function
    and instantiates the template with custom_kernel_myexp_float_float16_t_float16_t<float>.
    Template parameters can be mx.core.Dtype, int or bool.
  • Metal attributes used in source such as [[thread_position_in_grid]]
    These will be added as function arguments.
    All the attributes defined in Table 5.8 of the Metal Shading Language Specification <https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf>_ are supported.

Putting this all together, the generated function signature for myexp is as follows:

.. code-block:: cpp

template
[[kernel]] void custom_kernel_myexp_float_float16_t_float16_t(
const device float16_t* inp [[buffer(0)]],
device float16_t* out [[buffer(1)]],
uint3 thread_position_in_grid [[thread_position_in_grid]]) {

text
uint elem = thread_position_in_grid.x;
      T tmp = inp[elem];
      out[elem] = metal::exp(tmp);

}

template [[host_name("custom_kernel_myexp_float_float16_t_float16_t")]] [[kernel]] decltype(custom_kernel_myexp_float_float16_t_float16_t) custom_kernel_myexp_float_float16_t_float16_t;

Note: grid and threadgroup are parameters to the Metal dispatchThreads <https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/2866532-dispatchthreads>_
function. This means we will launch mx.prod(grid) threads, subdivided into
threadgroup size threadgroups. For optimal performance, each thread group
dimension should be less than or equal to the corresponding grid dimension.

Passing verbose=True when calling the kernel returned by
:func:fast.metal_kernel will print the generated code for debugging purposes.

Math Mode

By default :func:fast.metal_kernel compiles kernels with
compile_options={"math_mode": "safe"} so special values follow IEEE
behavior, for example exp(-inf) == 0. This is important for kernels such as
masked softmax where causal or sliding-window masks depend on exponentiating
-inf.

If your kernel does not rely on these edge cases, you can opt in to less strict
math with compile_options={"math_mode": "relaxed"} or
compile_options={"math_mode": "fast"}:

.. code-block:: python

kernel = mx.fast.metal_kernel(
name="my_kernel",
input_names=["x"],
output_names=["y"],
source=source,
compile_options={"math_mode": "relaxed"},
)

Using Shape/Strides

:func:fast.metal_kernel supports an argument ensure_row_contiguous which
is True by default. This will copy the array inputs if needed
before the kernel is launched to ensure that the memory layout is row
contiguous. Generally this makes writing the kernel easier, since we don't
have to worry about gaps or the ordering of the dims when indexing.

If we want to avoid this copy, :func:fast.metal_kernel automatically passes
a_shape, a_strides and a_ndim for each input array a if any are
present in source. We can then use MLX's built in indexing utils to fetch
the right elements for each thread.

Let's convert myexp above to support arbitrarily strided arrays without
relying on a copy from ensure_row_contiguous:

.. code-block:: python

source = """
uint elem = thread_position_in_grid.x;
// Utils from mlx/backend/metal/kernels/utils.h are automatically included
uint loc = elem_to_loc(elem, inp_shape, inp_strides, inp_ndim);
T tmp = inp[loc];
// Output arrays are always row contiguous
out[elem] = metal::exp(tmp);
"""

kernel = mx.fast.metal_kernel(
name="myexp_strided",
input_names=["inp"],
output_names=["out"],
source=source,
ensure_row_contiguous=False,
)

def exp_elementwise(a: mx.array):
outputs = kernel(
inputs=[a],
template=[("T", mx.float32)],
grid=(a.size, 1, 1),
threadgroup=(256, 1, 1),
output_shapes=[a.shape],
output_dtypes=[a.dtype],
)
return outputs[0]

a = mx.random.normal(shape=(4, 16)).astype(mx.float16)

make non-contiguous

a = a[::2]
b = exp_elementwise(a)
assert mx.allclose(b, mx.exp(a))

Complex Example

Let's implement a more complex example: grid_sample in "bilinear" mode.

We'll start with the following MLX implementation using standard ops:

.. code-block:: python

def grid_sample_ref(x, grid):
N, H_in, W_in, _ = x.shape
ix = ((grid[..., 0] + 1) * W_in - 1) / 2
iy = ((grid[..., 1] + 1) * H_in - 1) / 2

text
ix_nw = mx.floor(ix).astype(mx.int32)
  iy_nw = mx.floor(iy).astype(mx.int32)

  ix_ne = ix_nw + 1
  iy_ne = iy_nw

  ix_sw = ix_nw
  iy_sw = iy_nw + 1

  ix_se = ix_nw + 1
  iy_se = iy_nw + 1

  nw = (ix_se - ix)    * (iy_se - iy)
  ne = (ix    - ix_sw) * (iy_sw - iy)
  sw = (ix_ne - ix)    * (iy    - iy_ne)
  se = (ix    - ix_nw) * (iy    - iy_nw)

  I_nw = x[mx.arange(N)[:, None, None], iy_nw, ix_nw, :]
  I_ne = x[mx.arange(N)[:, None, None], iy_ne, ix_ne, :]
  I_sw = x[mx.arange(N)[:, None, None], iy_sw, ix_sw, :]
  I_se = x[mx.arange(N)[:, None, None], iy_se, ix_se, :]

  mask_nw = (iy_nw >= 0) & (iy_nw <= H_in - 1) & (ix_nw >= 0) & (ix_nw <= W_in - 1)
  mask_ne = (iy_ne >= 0) & (iy_ne <= H_in - 1) & (ix_ne >= 0) & (ix_ne <= W_in - 1)
  mask_sw = (iy_sw >= 0) & (iy_sw <= H_in - 1) & (ix_sw >= 0) & (ix_sw <= W_in - 1)
  mask_se = (iy_se >= 0) & (iy_se <= H_in - 1) & (ix_se >= 0) & (ix_se <= W_in - 1)

  I_nw *= mask_nw[..., None]
  I_ne *= mask_ne[..., None]
  I_sw *= mask_sw[..., None]
  I_se *= mask_se[..., None]

  output = nw[..., None] * I_nw + ne[..., None] * I_ne + sw[..., None] * I_sw + se[..., None] * I_se

  return output

Now let's use :func:custom_function together with :func:fast.metal_kernel
to write a fast GPU kernel for both the forward and backward passes.

First we'll implement the forward pass as a fused kernel:

.. code-block:: python

source = """
uint elem = thread_position_in_grid.x;
int H = x_shape[1];
int W = x_shape[2];
int C = x_shape[3];
int gH = grid_shape[1];
int gW = grid_shape[2];

text
int w_stride = C;
  int h_stride = W * w_stride;
  int b_stride = H * h_stride;

  uint grid_idx = elem / C * 2;
  float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
  float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;

  int ix_nw = floor(ix);
  int iy_nw = floor(iy);

  int ix_ne = ix_nw + 1;
  int iy_ne = iy_nw;

  int ix_sw = ix_nw;
  int iy_sw = iy_nw + 1;

  int ix_se = ix_nw + 1;
  int iy_se = iy_nw + 1;

  T nw = (ix_se - ix)    * (iy_se - iy);
  T ne = (ix    - ix_sw) * (iy_sw - iy);
  T sw = (ix_ne - ix)    * (iy    - iy_ne);
  T se = (ix    - ix_nw) * (iy    - iy_nw);

  int batch_idx = elem / C / gH / gW * b_stride;
  int channel_idx = elem % C;
  int base_idx = batch_idx + channel_idx;

  T I_nw = x[base_idx + iy_nw * h_stride + ix_nw * w_stride];
  T I_ne = x[base_idx + iy_ne * h_stride + ix_ne * w_stride];
  T I_sw = x[base_idx + iy_sw * h_stride + ix_sw * w_stride];
  T I_se = x[base_idx + iy_se * h_stride + ix_se * w_stride];

  I_nw = iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1 ? I_nw : 0;
  I_ne = iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1 ? I_ne : 0;
  I_sw = iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1 ? I_sw : 0;
  I_se = iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1 ? I_se : 0;

  out[elem] = nw * I_nw + ne * I_ne + sw * I_sw + se * I_se;

"""

kernel = mx.fast.metal_kernel(
name="grid_sample",
input_names=["x", "grid"],
output_names=["out"],
source=source,
)

@mx.custom_function
def grid_sample(x, grid):

text
assert x.ndim == 4, "`x` must be 4D."
  assert grid.ndim == 4, "`grid` must be 4D."

  B, _, _, C = x.shape
  _, gN, gM, D = grid.shape
  out_shape = (B, gN, gM, C)

  assert D == 2, "Last dim of `grid` must be size 2."

  outputs = kernel(
      inputs=[x, grid],
      template=[("T", x.dtype)],
      output_shapes=[out_shape],
      output_dtypes=[x.dtype],
      grid=(np.prod(out_shape), 1, 1),
      threadgroup=(256, 1, 1),
  )
  return outputs[0]

For a reasonably sized input such as:

.. code-block:: python

x.shape = (8, 1024, 1024, 64)
grid.shape = (8, 256, 256, 2)

On an M1 Max, we see a big performance improvement:

55.7ms -> 6.7ms => 8x speed up

Grid Sample VJP

Since we decorated grid_sample with :func:custom_function, we can now
define its custom vjp transform so MLX can differentiate it.

The backwards pass requires atomically updating x_grad/grid_grad and so
requires a few extra :func:fast.metal_kernel features:

  • init_value=0
    Initialize all of the kernel's outputs to this value before it runs. This allows us to update only part of the output arrays with the kernel.

  • atomic_outputs=True
    Designate all of the kernel outputs as atomic in the function signature.
    This means we can use Metal's atomic features to simultaneously update the x_grad and grid_grad arrays from multiple threadgroups.
    See section 6.15 of the Metal Shading Language Specification <https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf>_ for more details.

We can then implement the backwards pass as follows:

.. code-block:: python

source = """
uint elem = thread_position_in_grid.x;
int H = x_shape[1];
int W = x_shape[2];
int C = x_shape[3];
// Pad C to the nearest larger simdgroup size multiple
int C_padded = ceildiv(C, threads_per_simdgroup) * threads_per_simdgroup;

text
int gH = grid_shape[1];
  int gW = grid_shape[2];

  int w_stride = C;
  int h_stride = W * w_stride;
  int b_stride = H * h_stride;

  uint grid_idx = elem / C_padded * 2;
  float ix = ((grid[grid_idx] + 1) * W - 1) / 2;
  float iy = ((grid[grid_idx + 1] + 1) * H - 1) / 2;

  int ix_nw = floor(ix);
  int iy_nw = floor(iy);

  int ix_ne = ix_nw + 1;
  int iy_ne = iy_nw;

  int ix_sw = ix_nw;
  int iy_sw = iy_nw + 1;

  int ix_se = ix_nw + 1;
  int iy_se = iy_nw + 1;

  T nw = (ix_se - ix)    * (iy_se - iy);
  T ne = (ix    - ix_sw) * (iy_sw - iy);
  T sw = (ix_ne - ix)    * (iy    - iy_ne);
  T se = (ix    - ix_nw) * (iy    - iy_nw);

  int batch_idx = elem / C_padded / gH / gW * b_stride;
  int channel_idx = elem % C_padded;
  int base_idx = batch_idx + channel_idx;

  T gix = T(0);
  T giy = T(0);
  if (channel_idx < C) {
      int cot_index = elem / C_padded * C + channel_idx;
      T cot = cotangent[cot_index];
      if (iy_nw >= 0 && iy_nw <= H - 1 && ix_nw >= 0 && ix_nw <= W - 1) {
          int offset = base_idx + iy_nw * h_stride + ix_nw * w_stride;
          atomic_fetch_add_explicit(&x_grad[offset], nw * cot, memory_order_relaxed);

          T I_nw = x[offset];
          gix -= I_nw * (iy_se - iy) * cot;
          giy -= I_nw * (ix_se - ix) * cot;
      }
      if (iy_ne >= 0 && iy_ne <= H - 1 && ix_ne >= 0 && ix_ne <= W - 1) {
          int offset = base_idx + iy_ne * h_stride + ix_ne * w_stride;
          atomic_fetch_add_explicit(&x_grad[offset], ne * cot, memory_order_relaxed);

          T I_ne = x[offset];
          gix += I_ne * (iy_sw - iy) * cot;
          giy -= I_ne * (ix - ix_sw) * cot;
      }
      if (iy_sw >= 0 && iy_sw <= H - 1 && ix_sw >= 0 && ix_sw <= W - 1) {
          int offset = base_idx + iy_sw * h_stride + ix_sw * w_stride;
          atomic_fetch_add_explicit(&x_grad[offset], sw * cot, memory_order_relaxed);

          T I_sw = x[offset];
          gix -= I_sw * (iy - iy_ne) * cot;
          giy += I_sw * (ix_ne - ix) * cot;
      }
      if (iy_se >= 0 && iy_se <= H - 1 && ix_se >= 0 && ix_se <= W - 1) {
          int offset = base_idx + iy_se * h_stride + ix_se * w_stride;
          atomic_fetch_add_explicit(&x_grad[offset], se * cot, memory_order_relaxed);

          T I_se = x[offset];
          gix += I_se * (iy - iy_nw) * cot;
          giy += I_se * (ix - ix_nw) * cot;
      }
  }

  T gix_mult = W / 2;
  T giy_mult = H / 2;

  // Reduce across each simdgroup first.
  // This is much faster than relying purely on atomics.
  gix = simd_sum(gix);
  giy = simd_sum(giy);

  if (thread_index_in_simdgroup == 0) {
      atomic_fetch_add_explicit(&grid_grad[grid_idx], gix * gix_mult, memory_order_relaxed);
      atomic_fetch_add_explicit(&grid_grad[grid_idx + 1], giy * giy_mult, memory_order_relaxed);
  }

"""
kernel = mx.fast.metal_kernel(
name="grid_sample_grad",
input_names=["x", "grid", "cotangent"],
output_names=["x_grad", "grid_grad"],
source=source,
atomic_outputs=True,
)

@grid_sample.vjp
def grid_sample_vjp(primals, cotangent, _):
x, grid = primals
B, _, _, C = x.shape
_, gN, gM, D = grid.shape

text
assert D == 2, "Last dim of `grid` must be size 2."

  # pad the output channels to simd group size
  # so that our `simd_sum`s don't overlap.
  simdgroup_size = 32
  C_padded = (C + simdgroup_size - 1) // simdgroup_size * simdgroup_size
  grid_size = B * gN * gM * C_padded
  outputs = kernel(
      inputs=[x, grid, cotangent],
      template=[("T", x.dtype)],
      output_shapes=[x.shape, grid.shape],
      output_dtypes=[x.dtype, x.dtype],
      grid=(grid_size, 1, 1),
      threadgroup=(256, 1, 1),
      init_value=0,
  )
  return outputs[0], outputs[1]

There's an even larger speed up for the vjp:

676.4ms -> 16.7ms => 40x speed up

File: docs/src/dev/extensions.rst

Custom Extensions in MLX

You can extend MLX with custom operations on the CPU or GPU. This guide
explains how to do that with a simple example.

Introducing the Example

Let's say you would like an operation that takes in two arrays, x and
y, scales them both by coefficients alpha and beta respectively,
and then adds them together to get the result z = alpha * x + beta * y.
You can do that in MLX directly:

.. code-block:: python

text
import mlx.core as mx

def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
    return alpha * x + beta * y

This function performs that operation while leaving the implementation and
function transformations to MLX.

However, you may want to customize the underlying implementation, perhaps to
make it faster. In this tutorial we will go through adding custom extensions.
It will cover:

  • The structure of the MLX library.
  • Implementing a CPU operation.
  • Implementing a GPU operation using metal.
  • Adding the vjp and jvp function transformation.
  • Building a custom extension and binding it to python.

Operations and Primitives

Operations in MLX build the computation graph. Primitives provide the rules for
evaluating and transforming the graph. Let's start by discussing operations in
more detail.

Operations
^^^^^^^^^^^

Operations are the front-end functions that operate on arrays. They are defined
in the C++ API (:ref:cpp_ops), and the Python API (:ref:ops) binds them.

We would like an operation :meth:axpby that takes in two arrays, x and
y, and two scalars, alpha and beta. This is how to define it in
C++:

.. code-block:: C++

text
/**
*  Scale and sum two vectors element-wise
*  z = alpha * x + beta * y
*
*  Use NumPy-style broadcasting between x and y
*  Inputs are upcasted to floats if needed
**/
array axpby(
    const array& x, // Input array x
    const array& y, // Input array y
    const float alpha, // Scaling factor for x
    const float beta, // Scaling factor for y
    StreamOrDevice s = {} // Stream on which to schedule the operation
);

The simplest way to implement this is with existing operations:

.. code-block:: C++

text
array axpby(
    const array& x, // Input array x
    const array& y, // Input array y
    const float alpha, // Scaling factor for x
    const float beta, // Scaling factor for y
    StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
) {
    // Scale x and y on the provided stream
    auto ax = multiply(array(alpha), x, s);
    auto by = multiply(array(beta), y, s);

    // Add and return
    return add(ax, by, s);
}

The operations themselves do not contain the implementations that act on the
data, nor do they contain the rules of transformations. Rather, they are an
easy to use interface that use :class:Primitive building blocks.

Primitives
^^^^^^^^^^^

A :class:Primitive is part of the computation graph of an :class:array. It
defines how to create output arrays given input arrays. Further, a
:class:Primitive has methods to run on the CPU or GPU and for function
transformations such as vjp and jvp. Let's go back to our example to be
more concrete:

.. code-block:: C++

text
class Axpby : public Primitive {
  public:
    explicit Axpby(Stream stream, float alpha, float beta)
        : Primitive(stream), alpha_(alpha), beta_(beta){};

    /**
    * A primitive must know how to evaluate itself on the CPU/GPU
    * for the given inputs and populate the output array.
    *
    * To avoid unnecessary allocations, the evaluation function
    * is responsible for allocating space for the array.
    */
    void eval_cpu(
        const std::vector<array>& inputs,
        std::vector<array>& outputs) override;
    void eval_gpu(
        const std::vector<array>& inputs,
        std::vector<array>& outputs) override;

    /** The Jacobian-vector product. */
    std::vector<array> jvp(
        const std::vector<array>& primals,
        const std::vector<array>& tangents,
        const std::vector& argnums) override;

    /** The vector-Jacobian product. */
    std::vector<array> vjp(
        const std::vector<array>& primals,
        const std::vector<array>& cotangents,
        const std::vector& argnums,
        const std::vector<array>& outputs) override;

    /**
    * The primitive must know how to vectorize itself across
    * the given axes. The output is a pair containing the array
    * representing the vectorized computation and the axis which
    * corresponds to the output vectorized dimension.
    */
    std::pair<std::vector<array>, std::vector> vmap(
        const std::vector<array>& inputs,
        const std::vector& axes) override;

    /** The name of primitive. */
    const char* name() const override {
      return "Axpby";
    }

    /** Equivalence check **/
    bool is_equivalent(const Primitive& other) const override;

  private:
    float alpha_;
    float beta_;
};

The :class:Axpby class derives from the base :class:Primitive class. The
:class:Axpby treats alpha and beta as parameters. It then provides
implementations of how the output array is produced given the inputs through
:meth:Axpby::eval_cpu and :meth:Axpby::eval_gpu. It also provides rules
of transformations in :meth:Axpby::jvp, :meth:Axpby::vjp, and
:meth:Axpby::vmap.

Using the Primitive
^^^^^^^^^^^^^^^^^^^

Operations can use this :class:Primitive to add a new :class:array to the
computation graph. An :class:array can be constructed by providing its data
type, shape, the :class:Primitive that computes it, and the :class:array
inputs that are passed to the primitive.

Let's reimplement our operation now in terms of our :class:Axpby primitive.

.. code-block:: C++

text
array axpby(
    const array& x, // Input array x
    const array& y, // Input array y
    const float alpha, // Scaling factor for x
    const float beta, // Scaling factor for y
    StreamOrDevice s /* = {} */ // Stream on which to schedule the operation
) {
    // Promote dtypes between x and y as needed
    auto promoted_dtype = promote_types(x.dtype(), y.dtype());

    // Upcast to float32 for non-floating point inputs x and y
    auto out_dtype = issubdtype(promoted_dtype, float32)
        ? promoted_dtype
        : promote_types(promoted_dtype, float32);

    // Cast x and y up to the determined dtype (on the same stream s)
    auto x_casted = astype(x, out_dtype, s);
    auto y_casted = astype(y, out_dtype, s);

    // Broadcast the shapes of x and y (on the same stream s)
    auto broadcasted_inputs = broadcast_arrays({x_casted, y_casted}, s);
    auto out_shape = broadcasted_inputs[0].shape();

    // Construct the array as the output of the Axpby primitive
    // with the broadcasted and upcasted arrays as inputs
    return array(
        /* const std::vector& shape = */ out_shape,
        /* Dtype dtype = */ out_dtype,
        /* std::unique_ptr primitive = */
        std::make_shared<Axpby>(to_stream(s), alpha, beta),
        /* const std::vector<array>& inputs = */ broadcasted_inputs);
}

This operation now handles the following:

#. Upcast inputs and resolve the output data type.
#. Broadcast the inputs and resolve the output shape.
#. Construct the primitive :class:Axpby using the given stream, alpha, and beta.
#. Construct the output :class:array using the primitive and the inputs.

Implementing the Primitive

No computation happens when we call the operation alone. The operation only
builds the computation graph. When we evaluate the output array, MLX schedules
the execution of the computation graph, and calls :meth:Axpby::eval_cpu or
:meth:Axpby::eval_gpu depending on the stream/device specified by the user.

.. warning::
When :meth:Primitive::eval_cpu or :meth:Primitive::eval_gpu are called,
no memory has been allocated for the output array. It falls on the implementation
of these functions to allocate memory as needed.

Implementing the CPU Back-end
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Let's start by implementing :meth:Axpby::eval_cpu.

The method will go over each element of the output array, find the
corresponding input elements of x and y and perform the operation
point-wise. This is captured in the templated function :meth:axpby_impl.

.. code-block:: C++

template
void axpby_impl(
const mx::array& x,
const mx::array& y,
mx::array& out,
float alpha_,
float beta_,
mx::Stream stream) {
out.set_data(mx::allocator::malloc(out.nbytes()));

text
// Get the CPU command encoder and register input and output arrays
auto& encoder = mx::cpu::get_command_encoder(stream);
encoder.set_input_array(x);
encoder.set_input_array(y);
encoder.set_output_array(out);

// Launch the CPU kernel
encoder.dispatch([x_ptr = x.data<T>(),
                  y_ptr = y.data<T>(),
                  out_ptr = out.data<T>(),
                  size = out.size(),
                  shape = out.shape(),
                  x_strides = x.strides(),
                  y_strides = y.strides(),
                  alpha_,
                  beta_]() {

  // Cast alpha and beta to the relevant types
  T alpha = static_cast<T>(alpha_);
  T beta = static_cast<T>(beta_);

  // Do the element-wise operation for each output
  for (size_t out_idx = 0; out_idx < size; out_idx++) {
    // Map linear indices to offsets in x and y
    auto x_offset = mx::elem_to_loc(out_idx, shape, x_strides);
    auto y_offset = mx::elem_to_loc(out_idx, shape, y_strides);

    // We allocate the output to be contiguous and regularly strided
    // (defaults to row major) and hence it doesn't need additional mapping
    out_ptr[out_idx] = alpha * x_ptr[x_offset] + beta * y_ptr[y_offset];
  }
});

}

Our implementation should work for all incoming floating point arrays.
Accordingly, we add dispatches for float32, float16, bfloat16 and
complex64. We throw an error if we encounter an unexpected type.

.. code-block:: C++

text
void Axpby::eval_cpu(
    const std::vector<mx::array>& inputs,
    std::vector<mx::array>& outputs) {
  auto& x = inputs[0];
  auto& y = inputs[1];
  auto& out = outputs[0];

  // Dispatch to the correct dtype
  if (out.dtype() == mx::float32) {
    return axpby_impl<float>(x, y, out, alpha_, beta_, stream());
  } else if (out.dtype() == mx::float16) {
    return axpby_impl<mx::float16_t>(x, y, out, alpha_, beta_, stream());
  } else if (out.dtype() == mx::bfloat16) {
    return axpby_impl<mx::bfloat16_t>(x, y, out, alpha_, beta_, stream());
  } else if (out.dtype() == mx::complex64) {
    return axpby_impl<mx::complex64_t>(x, y, out, alpha_, beta_, stream());
  } else {
    throw std::runtime_error(
        "Axpby is only supported for floating point types.");
  }
}

Just this much is enough to run the operation :meth:axpby on a CPU stream! If
you do not plan on running the operation on the GPU or using transforms on
computation graphs that contain :class:Axpby, you can stop implementing the
primitive here.

Implementing the GPU Back-end
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Apple silicon devices address their GPUs using the Metal_ shading language, and
GPU kernels in MLX are written using Metal.

.. note::

text
Here are some helpful resources if you are new to Metal:

* A walkthrough of the metal compute pipeline: `Metal Example`_
* Documentation for metal shading language: `Metal Specification`_
* Using metal from C++: `Metal-cpp`_

.. note::

text
As of Metal 4.1, the implicit address space of ``this`` in a member function
is ``__metal_generic`` rather than ``thread``. Member functions in Metal
shading code should therefore carry an explicit address space qualifier,
written after the parameter list:

.. code-block:: C++

    struct Accumulator {
        float vals[4];
        thread float& at(short i) thread {
            return vals[i];
        }
    };

Without the trailing ``thread``, returning a ``thread``-qualified reference
to a member no longer compiles::

    error: reference to type 'thread float' could not bind to an lvalue
    of type '__metal_generic float'

Metal libraries built with ``mlx_build_metallib`` are compiled with
``-Wmetal-addr-spaces``, which reports a missing qualifier as a warning on
toolchains where it is not yet an error.

Let's keep the GPU kernel simple. We will launch exactly as many threads as
there are elements in the output. Each thread will pick the element it needs
from x and y, do the point-wise operation, and update its assigned
element in the output.

.. code-block:: C++

text
template <typename T>
[[kernel]] void axpby_general(
        device const T* x [[buffer(0)]],
        device const T* y [[buffer(1)]],
        device T* out [[buffer(2)]],
        constant const float& alpha [[buffer(3)]],
        constant const float& beta [[buffer(4)]],
        constant const int* shape [[buffer(5)]],
        constant const int64_t* x_strides [[buffer(6)]],
        constant const int64_t* y_strides [[buffer(7)]],
        constant const int& ndim [[buffer(8)]],
        uint index [[thread_position_in_grid]]) {
    // Convert linear indices to offsets in array
    auto x_offset = elem_to_loc(index, shape, x_strides, ndim);
    auto y_offset = elem_to_loc(index, shape, y_strides, ndim);

    // Do the operation and update the output
    out[index] =
        static_cast<T>(alpha) * x[x_offset] + static_cast<T>(beta) * y[y_offset];
}

We then need to instantiate this template for all floating point types and give
each instantiation a unique host name so we can identify it.

.. code-block:: C++

text
instantiate_kernel("axpby_general_float32", axpby_general, float)
instantiate_kernel("axpby_general_float16", axpby_general, float16_t)
instantiate_kernel("axpby_general_bfloat16", axpby_general, bfloat16_t)
instantiate_kernel("axpby_general_complex64", axpby_general, complex64_t)

The logic to determine the kernel, set the inputs, resolve the grid dimensions,
and dispatch to the GPU are contained in :meth:Axpby::eval_gpu as shown
below.

.. code-block:: C++

text
/** Evaluate primitive on GPU */
void Axpby::eval_gpu(
  const std::vector<array>& inputs,
  std::vector<array>& outputs) {
    // Prepare inputs
    assert(inputs.size() == 2);
    auto& x = inputs[0];
    auto& y = inputs[1];
    auto& out = outputs[0];

    // Each primitive carries the stream it should execute on
    // and each stream carries its device identifiers
    auto& s = stream();
    // We get the needed metal device using the stream
    auto& d = metal::device(s.device);

    // Allocate output memory
    out.set_data(allocator::malloc(out.nbytes()));

    // Resolve name of kernel
    std::stream kname;
    kname = "axpby_general_" + type_to_name(out);

    // Load the metal library
    auto lib = d.get_library("mlx_ext", current_binary_dir());

    // Make a kernel from this metal library
    auto kernel = d.get_kernel(kname, lib);

    // Prepare to encode kernel
    auto& compute_encoder = mx::metal::get_command_encoder(s);
    compute_encoder.set_compute_pipeline_state(kernel);

    // Kernel parameters are registered with buffer indices corresponding to
    // those in the kernel declaration at axpby.metal
    int ndim = out.ndim();
    size_t nelem = out.size();

    // Encode input arrays to kernel
    compute_encoder.set_input_array(x, 0);
    compute_encoder.set_input_array(y, 1);

    // Encode output arrays to kernel
    compute_encoder.set_output_array(out, 2);

    // Encode alpha and beta
    compute_encoder.set_bytes(alpha_, 3);
    compute_encoder.set_bytes(beta_, 4);

    // Encode shape, strides and ndim
    compute_encoder.set_vector_bytes(x.shape(), 5);
    compute_encoder.set_vector_bytes(x.strides(), 6);
    compute_encoder.set_bytes(y.strides(), 7);
    compute_encoder.set_bytes(ndim, 8);

    // We launch 1 thread for each input and make sure that the number of
    // threads in any given threadgroup is not higher than the max allowed
    size_t tgp_size = std::min(nelem, kernel->maxTotalThreadsPerThreadgroup());

    // Fix the 3D size of each threadgroup (in terms of threads)
    MTL::Size group_dims = MTL::Size(tgp_size, 1, 1);

    // Fix the 3D size of the launch grid (in terms of threads)
    MTL::Size grid_dims = MTL::Size(nelem, 1, 1);

    // Launch the grid with the given number of threads divided among
    // the given threadgroups
    compute_encoder.dispatch_threads(grid_dims, group_dims);
}

We can now call the :meth:axpby operation on both the CPU and the GPU!

A few things to note about MLX and Metal before moving on. MLX keeps track of
the active command_buffer and the MTLCommandBuffer to which it is
associated. We rely on :meth:metal::get_command_encoder to give us the active
metal compute command encoder instead of building a new one and calling
:meth:compute_encoder->end_encoding at the end. MLX adds kernels (compute
pipelines) to the active command buffer until some specified limit is hit or
the command buffer needs to be flushed for synchronization.

Primitive Transforms
^^^^^^^^^^^^^^^^^^^^^

Next, let's add implementations for transformations in a :class:Primitive.
These transformations can be built on top of other operations, including the
one we just defined:

.. code-block:: C++

text
/** The Jacobian-vector product. */
std::vector<array> Axpby::jvp(
        const std::vector<array>& primals,
        const std::vector<array>& tangents,
        const std::vector& argnums) {
    // Forward mode diff that pushes along the tangents
    // The jvp transform on the primitive can be built with ops
    // that are scheduled on the same stream as the primitive

    // If argnums = {0}, we only push along x in which case the
    // jvp is just the tangent scaled by alpha
    // Similarly, if argnums = {1}, the jvp is just the tangent
    // scaled by beta
    if (argnums.size() > 1) {
        auto scale = argnums[0] == 0 ? alpha_ : beta_;
        auto scale_arr = array(scale, tangents[0].dtype());
        return {multiply(scale_arr, tangents[0], stream())};
    }
    // If argnums = {0, 1}, we take contributions from both
    // which gives us jvp = tangent_x * alpha + tangent_y * beta
    else {
        return {axpby(tangents[0], tangents[1], alpha_, beta_, stream())};
    }
}

.. code-block:: C++

text
/** The vector-Jacobian product. */
std::vector<array> Axpby::vjp(
        const std::vector<array>& primals,
        const std::vector<array>& cotangents,
        const std::vector& argnums,
        const std::vector& /* unused */) {
    // Reverse mode diff
    std::vector<array> vjps;
    for (auto arg : argnums) {
        auto scale = arg == 0 ? alpha_ : beta_;
        auto scale_arr = array(scale, cotangents[0].dtype());
        vjps.push_back(multiply(scale_arr, cotangents[0], stream()));
    }
    return vjps;
}

Note, a transformation does not need to be fully defined to start using
the :class:Primitive.

.. code-block:: C++

text
/** Vectorize primitive along given axis */
std::pair<std::vector<array>, std::vector> Axpby::vmap(
        const std::vector<array>& inputs,
        const std::vector& axes) {
    throw std::runtime_error("[Axpby] vmap not implemented.");
}

Building and Binding

Let's look at the overall directory structure first.

| extensions
| ├── axpby
| │ ├── axpby.cpp
| │ ├── axpby.h
| │ └── axpby.metal
| ├── mlx_sample_extensions
| │ └── init.py
| ├── bindings.cpp
| ├── CMakeLists.txt
| └── setup.py

  • extensions/axpby/ defines the C++ extension library
  • extensions/mlx_sample_extensions sets out the structure for the
    associated Python package
  • extensions/bindings.cpp provides Python bindings for our operation
  • extensions/CMakeLists.txt holds CMake rules to build the library and
    Python bindings
  • extensions/setup.py holds the setuptools rules to build and install
    the Python package

Binding to Python
^^^^^^^^^^^^^^^^^^

We use nanobind_ to build a Python API for the C++ library. Since bindings for
components such as :class:mlx.core.array, :class:mlx.core.stream, etc. are
already provided, adding our :meth:axpby is simple.

.. code-block:: C++

NB_MODULE(_ext, m) {
m.doc() = "Sample extension for MLX";

text
m.def(
        "axpby",
        &axpby,
        "x"_a,
        "y"_a,
        "alpha"_a,
        "beta"_a,
        nb::kw_only(),
        "stream"_a = nb::none(),
        R"(
            Scale and sum two vectors element-wise
            ``z = alpha * x + beta * y``

            Follows numpy style broadcasting between ``x`` and ``y``
            Inputs are upcasted to floats if needed

            Args:
                x (array): Input array.
                y (array): Input array.
                alpha (float): Scaling factor for ``x``.
                beta (float): Scaling factor for ``y``.

            Returns:
                array: ``alpha * x + beta * y``
        )");
}

Most of the complexity in the above example comes from additional bells and
whistles such as the literal names and doc-strings.

.. warning::

text
:mod:`mlx.core` must be imported before importing
:mod:`mlx_sample_extensions` as defined by the nanobind module above to
ensure that the casters for :mod:`mlx.core` components like
:class:`mlx.core.array` are available.

.. _Building with CMake:

Building with CMake
^^^^^^^^^^^^^^^^^^^^

Building the C++ extension library only requires that you find_package(MLX CONFIG) and then link it to your library.

.. code-block:: cmake

text
# Add library
add_library(mlx_ext)

# Add sources
target_sources(
    mlx_ext
    PUBLIC
    ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.cpp
)

# Add include headers
target_include_directories(
    mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR}
)

# Link to mlx
target_link_libraries(mlx_ext PUBLIC mlx)

We also need to build the attached Metal library. For convenience, we provide a
:meth:mlx_build_metallib function that builds a .metallib target given
sources, headers, destinations, etc. (defined in cmake/extension.cmake and
automatically imported with MLX package).

Here is what that looks like in practice:

.. code-block:: cmake

text
# Build metallib
if(MLX_BUILD_METAL)

mlx_build_metallib(
    TARGET mlx_ext_metallib
    TITLE mlx_ext
    SOURCES ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
    INCLUDE_DIRS ${PROJECT_SOURCE_DIR} ${MLX_INCLUDE_DIRS}
    OUTPUT_DIRECTORY ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}
)

add_dependencies(
    mlx_ext
    mlx_ext_metallib
)

endif()

Finally, we build the nanobind_ bindings

.. code-block:: cmake

text
nanobind_add_module(
  _ext
  NB_STATIC STABLE_ABI LTO NOMINSIZE
  NB_DOMAIN mlx
  ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp
)
target_link_libraries(_ext PRIVATE mlx_ext)

if(BUILD_SHARED_LIBS)
  target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
endif()

Building with setuptools
^^^^^^^^^^^^^^^^^^^^^^^^^^^^

Once we have set out the CMake build rules as described above, we can use the
build utilities defined in :mod:mlx.extension:

.. code-block:: python

text
from mlx import extension
from setuptools import setup

if __name__ == "__main__":
    setup(
        name="mlx_sample_extensions",
        version="0.0.0",
        description="Sample C++ and Metal extensions for MLX primitives.",
        ext_modules=[extension.CMakeExtension("mlx_sample_extensions._ext")],
        cmdclass={"build_ext": extension.CMakeBuild},
        packages=["mlx_sample_extensions"],
        package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]},
        extras_require={"dev":[]},
        zip_safe=False,
        python_requires=">=3.8",
    )

.. note::
We treat extensions/mlx_sample_extensions as the package directory
even though it only contains a __init__.py to ensure the following:

text
* :mod:`mlx.core` must be imported before importing :mod:`_ext`
* The C++ extension library and the metal library are co-located with the python
  bindings and copied together if the package is installed

To build the package, first install the build dependencies with pip install -r requirements.txt. You can then build inplace for development using
python setup.py build_ext -j8 --inplace (in extensions/)

This results in the directory structure:

| extensions
| ├── mlx_sample_extensions
| │ ├── init.py
| │ ├── libmlx_ext.dylib # C++ extension library
| │ ├── mlx_ext.metallib # Metal library
| │ └── _ext.cpython-3x-darwin.so # Python Binding
| ...

When you try to install using the command python -m pip install . (in
extensions/), the package will be installed with the same structure as
extensions/mlx_sample_extensions and the C++ and Metal library will be
copied along with the Python binding since they are specified as
package_data.

Usage

After installing the extension as described above, you should be able to simply
import the Python package and play with it as you would any other MLX operation.

Let's look at a simple script and its results:

.. code-block:: python

text
import mlx.core as mx
from mlx_sample_extensions import axpby

a = mx.ones((3, 4))
b = mx.ones((3, 4))
c = axpby(a, b, 4.0, 2.0, stream=mx.cpu)

print(f"c shape: {c.shape}")
print(f"c dtype: {c.dtype}")
print(f"c is correct: {mx.all(c == 6.0).item()}")

Output:

.. code-block::

text
c shape: [3, 4]
c dtype: float32
c is correct: True

Results
^^^^^^^

Let's run a quick benchmark and see how our new axpby operation compares
with the naive :meth:simple_axpby we first defined.

.. code-block:: python

text
import mlx.core as mx
from mlx_sample_extensions import axpby
import time

def simple_axpby(x: mx.array, y: mx.array, alpha: float, beta: float) -> mx.array:
    return alpha * x + beta * y

M = 4096
N = 4096

x = mx.random.normal((M, N))
y = mx.random.normal((M, N))
alpha = 4.0
beta = 2.0

mx.eval(x, y)

def bench(f):
    # Warm up
    for i in range(5):
        z = f(x, y, alpha, beta)
        mx.eval(z)

    # Timed run
    s = time.perf_counter()
    for i in range(100):
        z = f(x, y, alpha, beta)
        mx.eval(z)
    e = time.perf_counter()
    return 1000 * (e - s) / 100

simple_time = bench(simple_axpby)
custom_time = bench(axpby)

print(f"Simple axpby: {simple_time:.3f} ms | Custom axpby: {custom_time:.3f} ms")

The results are Simple axpby: 1.559 ms | Custom axpby: 0.774 ms. We see
modest improvements right away!

This operation is now good to be used to build other operations, in
:class:mlx.nn.Module calls, and also as a part of graph transformations like
:meth:grad.

Scripts

.. admonition:: Download the code

The full example code is available in mlx <https://github.com/ml-explore/mlx/tree/main/examples/extensions/>_.

.. _Accelerate: https://developer.apple.com/documentation/accelerate/blas?language=objc
.. _Metal: https://developer.apple.com/documentation/metal?language=objc
.. _Metal-cpp: https://developer.apple.com/metal/cpp/
.. _Metal Specification: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
.. _Metal Example: https://developer.apple.com/documentation/metal/performing_calculations_on_a_gpu?language=objc
.. _nanobind: https://nanobind.readthedocs.io/en/latest/

File: docs/src/dev/metal_debugger.rst

Metal Debugger

.. currentmodule:: mlx.core

Profiling is a key step for performance optimization. You can build MLX with
the MLX_METAL_DEBUG option to improve the Metal debugging and
optimization workflow. The MLX_METAL_DEBUG debug option:

  • Records source during Metal compilation, for later inspection while
    debugging.
  • Labels Metal objects such as command queues, improving capture readability.

To build with debugging enabled in Python prepend
CMAKE_ARGS="-DMLX_METAL_DEBUG=ON" to the build call.

The :func:metal.start_capture function initiates a capture of all MLX GPU
work.

.. note::

To capture a GPU trace you must run the application with
MTL_CAPTURE_ENABLED=1.

.. code-block:: python

text
import mlx.core as mx

a = mx.random.uniform(shape=(512, 512))
b = mx.random.uniform(shape=(512, 512))
mx.eval(a, b)

trace_file = "mlx_trace.gputrace"

# Make sure to run with MTL_CAPTURE_ENABLED=1 and
# that the path trace_file does not already exist.
mx.metal.start_capture(trace_file)

for _ in range(10):
  mx.eval(mx.add(a, b))

mx.metal.stop_capture()

You can open and replay the GPU trace in Xcode. The Dependencies view
has a great overview of all operations. Checkout the Metal debugger documentation_ for more information.

.. image:: ../_static/metal_debugger/capture.png
:class: dark-light

Xcode Workflow

You can skip saving to a path by running within Xcode. First, generate an
Xcode project using CMake.

.. code-block::

text
mkdir build && cd build
cmake .. -DMLX_METAL_DEBUG=ON -G Xcode
open mlx.xcodeproj

Select the metal_capture example schema and run.

.. image:: ../_static/metal_debugger/schema.png
:class: dark-light

.. _Metal debugger documentation: https://developer.apple.com/documentation/xcode/metal-debugger

File: docs/src/dev/mlx_in_cpp.rst

.. _mlx_in_cpp:

Using MLX in C++

You can use MLX in a C++ project with CMake.

.. note::

This guide is based one the following example using MLX in C++ <https://github.com/ml-explore/mlx/tree/main/examples/cmake_project>_

First install MLX:

.. code-block:: bash

pip install -U mlx

You can also install the MLX Python package from source or just the C++
library. For more information see the :ref:documentation on installing MLX .

Next make an example program in example.cpp:

.. code-block:: C++

#include

#include "mlx/mlx.h"

namespace mx = mlx::core;

int main() {
auto x = mx::array({1, 2, 3});
auto y = mx::array({1, 2, 3});
std::cout << x + y << std::endl;
return 0;
}

The next step is to setup a CMake file in CMakeLists.txt:

.. code-block:: cmake

cmake_minimum_required(VERSION 3.27)

project(example LANGUAGES CXX)

set(CMAKE_CXX_STANDARD 20)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

Depending on how you installed MLX, you may need to tell CMake where to
find it.

If you installed MLX with Python, then add the following to the CMake file:

.. code-block:: cmake

find_package(
Python 3.9
COMPONENTS Interpreter Development.Module
REQUIRED)
execute_process(
COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
OUTPUT_STRIP_TRAILING_WHITESPACE
OUTPUT_VARIABLE MLX_ROOT)

If you installed the MLX C++ package to a system path, then CMake should be
able to find it. If you installed it to a non-standard location or CMake can't
find MLX then set MLX_ROOT to the location where MLX is installed:

.. code-block:: cmake

set(MLX_ROOT "/path/to/mlx/")

Next, instruct CMake to find MLX:

.. code-block:: cmake

find_package(MLX CONFIG REQUIRED)

Finally, add the example.cpp program as an executable and link MLX.

.. code-block:: cmake

add_executable(example example.cpp)
target_link_libraries(example PRIVATE mlx)

You can build the example with:

.. code-block:: bash

cmake -B build -DCMAKE_BUILD_TYPE=Release
cmake --build build

And run it with:

.. code-block:: bash

./build/example

Note find_package(MLX CONFIG REQUIRED) sets the following variables:

.. list-table:: Package Variables
:widths: 20 20
:header-rows: 1

    • Variable
    • Description
    • MLX_FOUND
    • True if MLX is found
    • MLX_INCLUDE_DIRS
    • Include directory
    • MLX_LIBRARIES
    • Libraries to link against
    • MLX_CXX_FLAGS
    • Additional compiler flags
    • MLX_BUILD_ACCELERATE
    • True if MLX was built with Accelerate
    • MLX_BUILD_METAL
    • True if MLX was built with Metal

Explore Related Technologies & Topics

View All 9,111+ Repositories →

Discover similar high-velocity repositories, agent skills, and OpenAPI specifications across the ecosystem.

#mlx Browse mlx catalog → ⚡ MCP Servers Model Context Protocol → 🧠 Agent Skills Prompt & rule files → 🔌 OpenAPI Specs REST & Swagger schemas →