## File: README.md # MLX [**Quickstart**](#quickstart) | [**Installation**](#installation) | [**Documentation**](https://ml-explore.github.io/mlx/build/html/index.html) | [**Examples**](#examples) [](https://circleci.com/gh/ml-explore/mlx) 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](https://github.com/ml-explore/mlx-c), and [Swift](https://github.com/ml-explore/mlx-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](https://numpy.org/doc/stable/index.html), [PyTorch](https://pytorch.org/), [Jax](https://github.com/google/jax), and [ArrayFire](https://arrayfire.org/). ## Examples The [MLX examples repo](https://github.com/ml-explore/mlx-examples) has a variety of examples, including: - [Transformer language model](https://github.com/ml-explore/mlx-examples/tree/main/transformer_lm) training. - Large-scale text generation with [LLaMA](https://github.com/ml-explore/mlx-examples/tree/main/llms/llama) and finetuning with [LoRA](https://github.com/ml-explore/mlx-examples/tree/main/lora). - Generating images with [Stable Diffusion](https://github.com/ml-explore/mlx-examples/tree/main/stable_diffusion). - Speech recognition with [OpenAI's Whisper](https://github.com/ml-explore/mlx-examples/tree/main/whisper). ## Quickstart See the [quick start guide](https://ml-explore.github.io/mlx/build/html/usage/quick_start.html) in the documentation. ## Installation MLX is available on [PyPI](https://pypi.org/project/mlx/). 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](https://ml-explore.github.io/mlx/build/html/install.html#) for more information on building the C++ and Python APIs from source. ## Contributing Check out the [contribution guidelines](https://github.com/ml-explore/mlx/tree/main/CONTRIBUTING.md) for more information on contributing to MLX. See the [docs](https://ml-explore.github.io/mlx/build/html/install.html) for more information on building from source, and running tests. We are grateful for all of [our contributors](https://github.com/ml-explore/mlx/tree/main/ACKNOWLEDGMENTS.md#Individual-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 `` to the function and instantiates the template with ``custom_kernel_myexp_float_float16_t_float16_t``. 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 `_ 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]]) { 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 `_ 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 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]; 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): 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 `_ 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; 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 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 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++ /** * 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++ 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++ 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& inputs, std::vector& outputs) override; void eval_gpu( const std::vector& inputs, std::vector& outputs) override; /** The Jacobian-vector product. */ std::vector jvp( const std::vector& primals, const std::vector& tangents, const std::vector& argnums) override; /** The vector-Jacobian product. */ std::vector vjp( const std::vector& primals, const std::vector& cotangents, const std::vector& argnums, const std::vector& 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> vmap( const std::vector& 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++ 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(to_stream(s), alpha, beta), /* const std::vector& 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())); // 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(), y_ptr = y.data(), out_ptr = out.data(), 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(alpha_); T beta = static_cast(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++ void Axpby::eval_cpu( const std::vector& inputs, std::vector& 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(x, y, out, alpha_, beta_, stream()); } else if (out.dtype() == mx::float16) { return axpby_impl(x, y, out, alpha_, beta_, stream()); } else if (out.dtype() == mx::bfloat16) { return axpby_impl(x, y, out, alpha_, beta_, stream()); } else if (out.dtype() == mx::complex64) { return axpby_impl(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:: 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:: 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++ template [[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(alpha) * x[x_offset] + static_cast(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++ 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++ /** Evaluate primitive on GPU */ void Axpby::eval_gpu( const std::vector& inputs, std::vector& 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++ /** The Jacobian-vector product. */ std::vector Axpby::jvp( const std::vector& primals, const std::vector& 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++ /** The vector-Jacobian product. */ std::vector Axpby::vjp( const std::vector& primals, const std::vector& cotangents, const std::vector& argnums, const std::vector& /* unused */) { // Reverse mode diff std::vector 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++ /** Vectorize primitive along given axis */ std::pair, std::vector> Axpby::vmap( const std::vector& 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"; 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:: :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 # 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 # 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 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 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: * :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 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:: 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 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 `_. .. _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 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:: 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++ `_ 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