mlx 1.0.0

data/mlx/docs/src/usage/unified_memory.rst

@@ -0,0 +1,78 @@
+.. _unified_memory:
+
+Unified Memory
+==============
+
+.. currentmodule:: mlx.core
+
+Apple silicon has a unified memory architecture. The CPU and GPU have direct
+access to the same memory pool. MLX is designed to take advantage of that.
+
+Concretely, when you make an array in MLX you don't have to specify its location:
+
+
+.. code-block:: python
+
+  a = mx.random.normal((100,))
+  b = mx.random.normal((100,))
+
+Both ``a`` and ``b`` live in unified memory.
+
+In MLX, rather than moving arrays to devices, you specify the device when you
+run the operation. Any device can perform any operation on ``a`` and ``b``
+without needing to move them from one memory location to another. For example:
+
+.. code-block:: python
+
+  mx.add(a, b, stream=mx.cpu)
+  mx.add(a, b, stream=mx.gpu)
+
+In the above, both the CPU and the GPU will perform the same add
+operation. The operations can (and likely will) be run in parallel since
+there are no dependencies between them. See :ref:`using_streams` for more
+information the semantics of streams in MLX.
+
+In the above ``add`` example, there are no dependencies between operations, so
+there is no possibility for race conditions. If there are dependencies, the
+MLX scheduler will automatically manage them. For example:
+
+.. code-block:: python
+
+  c = mx.add(a, b, stream=mx.cpu)
+  d = mx.add(a, c, stream=mx.gpu)
+
+In the above case, the second ``add`` runs on the GPU but it depends on the
+output of the first ``add`` which is running on the CPU. MLX will
+automatically insert a dependency between the two streams so that the second
+``add`` only starts executing after the first is complete and ``c`` is
+available.
+
+A Simple Example
+~~~~~~~~~~~~~~~~
+
+Here is a more interesting (albeit slightly contrived example) of how unified
+memory can be helpful. Suppose we have the following computation:
+
+.. code-block:: python
+
+  def fun(a, b, d1, d2):
+    x = mx.matmul(a, b, stream=d1)
+    for _ in range(500):
+        b = mx.exp(b, stream=d2)
+    return x, b
+
+which we want to run with the following arguments:
+
+.. code-block:: python
+
+  a = mx.random.uniform(shape=(4096, 512))
+  b = mx.random.uniform(shape=(512, 4))
+
+The first ``matmul`` operation is a good fit for the GPU since it's more
+compute dense. The second sequence of operations are a better fit for the CPU,
+since they are very small and would probably be overhead bound on the GPU.
+
+If we time the computation fully on the GPU, we get 2.8 milliseconds. But if we
+run the computation with ``d1=mx.gpu`` and ``d2=mx.cpu``, then the time is only
+about 1.4 milliseconds, about twice as fast. These times were measured on an M1
+Max.

data/mlx/docs/src/usage/using_streams.rst

@@ -0,0 +1,18 @@
+.. _using_streams:
+
+Using Streams
+=============
+
+.. currentmodule:: mlx.core
+
+Specifying the :obj:`Stream`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+All operations (including random number generation) take an optional
+keyword argument ``stream``. The ``stream`` kwarg specifies which
+:obj:`Stream` the operation should run on. If the stream is unspecified then
+the operation is run on the default stream of the default device:
+``mx.default_stream(mx.default_device())``.  The ``stream`` kwarg can also
+be a :obj:`Device` (e.g. ``stream=my_device``) in which case the operation is
+run on the default stream of the provided device
+``mx.default_stream(my_device)``.

data/mlx/examples/cmake_project/CMakeLists.txt

@@ -0,0 +1,22 @@
+cmake_minimum_required(VERSION 3.27)
+
+project(example LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+# Comment the following two commands only the MLX C++ library is installed and
+# set(MLX_ROOT "/path/to/mlx") directly if needed.
+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)
+
+find_package(MLX CONFIG REQUIRED)
+
+add_executable(example example.cpp)
+target_link_libraries(example PRIVATE mlx)

data/mlx/examples/cmake_project/README.md

@@ -0,0 +1,26 @@
+## Build and Run 
+
+Install MLX with Python:
+
+```bash
+pip install mlx>=0.22
+```
+
+Build the C++ example:
+
+```bash
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build
+```
+
+Run the C++ example:
+
+```
+./build/example
+```
+
+which should output:
+
+```
+array([2, 4, 6], dtype=int32)
+```

data/mlx/examples/cmake_project/example.cpp

@@ -0,0 +1,14 @@
+// Copyright © 2024 Apple Inc.
+
+#include <iostream>
+
+#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;
+}

data/mlx/examples/cpp/CMakeLists.txt

@@ -0,0 +1,12 @@
+function(build_example SRCFILE)
+  get_filename_component(src_name ${SRCFILE} NAME_WE)
+  set(target "${src_name}")
+  add_executable(${target} ${SRCFILE})
+  target_link_libraries(${target} PRIVATE mlx)
+endfunction(build_example)
+
+build_example(tutorial.cpp)
+build_example(linear_regression.cpp)
+build_example(logistic_regression.cpp)
+build_example(metal_capture.cpp)
+build_example(distributed.cpp)

data/mlx/examples/cpp/distributed.cpp

@@ -0,0 +1,22 @@
+// Copyright © 2024 Apple Inc.
+
+#include <iostream>
+
+#include "mlx/mlx.h"
+
+namespace mx = mlx::core;
+
+int main() {
+  if (!mx::distributed::is_available()) {
+    std::cout << "No communication backend found" << std::endl;
+    return 1;
+  }
+
+  auto global_group = mx::distributed::init();
+  std::cout << global_group.rank() << " / " << global_group.size() << std::endl;
+
+  mx::array x = mx::ones({10});
+  mx::array out = mx::distributed::all_sum(x, global_group);
+
+  std::cout << out << std::endl;
+}

data/mlx/examples/cpp/linear_regression.cpp

@@ -0,0 +1,54 @@
+// Copyright © 2023 Apple Inc.
+
+#include <chrono>
+#include <cmath>
+#include <iostream>
+
+#include "mlx/mlx.h"
+#include "timer.h"
+
+/**
+ * An example of linear regression with MLX.
+ */
+namespace mx = mlx::core;
+
+int main() {
+  int num_features = 100;
+  int num_examples = 1'000;
+  int num_iters = 10'000;
+  float learning_rate = 0.01;
+
+  // True parameters
+  auto w_star = mx::random::normal({num_features});
+
+  // The input examples (design matrix)
+  auto X = mx::random::normal({num_examples, num_features});
+
+  // Noisy labels
+  auto eps = 1e-2 * mx::random::normal({num_examples});
+  auto y = mx::matmul(X, w_star) + eps;
+
+  // Initialize random parameters
+  mx::array w = 1e-2 * mx::random::normal({num_features});
+
+  auto loss_fn = [&](mx::array w) {
+    auto yhat = mx::matmul(X, w);
+    return (0.5f / num_examples) * mx::sum(mx::square(yhat - y));
+  };
+
+  auto grad_fn = mx::grad(loss_fn);
+
+  auto tic = timer::time();
+  for (int it = 0; it < num_iters; ++it) {
+    auto grads = grad_fn(w);
+    w = w - learning_rate * grads;
+    mx::eval(w);
+  }
+  auto toc = timer::time();
+
+  auto loss = loss_fn(w);
+  auto error_norm = std::sqrt(mx::sum(mx::square(w - w_star)).item<float>());
+  auto throughput = num_iters / timer::seconds(toc - tic);
+  std::cout << "Loss " << loss << ", |w - w*| = " << error_norm
+            << ", Throughput " << throughput << " (it/s)." << std::endl;
+}

data/mlx/examples/cpp/logistic_regression.cpp

@@ -0,0 +1,54 @@
+// Copyright © 2023 Apple Inc.
+
+#include <chrono>
+#include <cmath>
+#include <iostream>
+
+#include "mlx/mlx.h"
+#include "timer.h"
+
+/**
+ * An example of logistic regression with MLX.
+ */
+namespace mx = mlx::core;
+
+int main() {
+  int num_features = 100;
+  int num_examples = 1'000;
+  int num_iters = 10'000;
+  float learning_rate = 0.1;
+
+  // True parameters
+  auto w_star = mx::random::normal({num_features});
+
+  // The input examples
+  auto X = mx::random::normal({num_examples, num_features});
+
+  // Labels
+  auto y = mx::matmul(X, w_star) > 0;
+
+  // Initialize random parameters
+  mx::array w = 1e-2 * mx::random::normal({num_features});
+
+  auto loss_fn = [&](mx::array w) {
+    auto logits = mx::matmul(X, w);
+    auto scale = (1.0f / num_examples);
+    return scale * mx::sum(mx::logaddexp(mx::array(0.0f), logits) - y * logits);
+  };
+
+  auto grad_fn = mx::grad(loss_fn);
+
+  auto tic = timer::time();
+  for (int it = 0; it < num_iters; ++it) {
+    auto grads = grad_fn(w);
+    w = w - learning_rate * grads;
+    mx::eval(w);
+  }
+  auto toc = timer::time();
+
+  auto loss = loss_fn(w);
+  auto acc = mx::sum((mx::matmul(X, w) > 0) == y) / num_examples;
+  auto throughput = num_iters / timer::seconds(toc - tic);
+  std::cout << "Loss " << loss << ", Accuracy, " << acc << ", Throughput "
+            << throughput << " (it/s)." << std::endl;
+}

data/mlx/examples/cpp/metal_capture.cpp

@@ -0,0 +1,31 @@
+// Copyright © 2024 Apple Inc.
+
+#include <cassert>
+#include <iostream>
+
+#include "mlx/mlx.h"
+
+namespace mx = mlx::core;
+
+int main() {
+  // To use Metal debugging and profiling:
+  // 1. Build with the MLX_METAL_DEBUG CMake option (i.e. -DMLX_METAL_DEBUG=ON).
+  // 2. Run with MTL_CAPTURE_ENABLED=1.
+  mx::metal::start_capture("mlx_trace.gputrace");
+
+  // Start at index two because the default GPU and CPU streams have indices
+  // zero and one, respectively. This naming matches the label assigned to each
+  // stream's command queue.
+  auto s2 = new_stream(mx::Device::gpu);
+  auto s3 = new_stream(mx::Device::gpu);
+
+  auto a = mx::arange(1.f, 10.f, 1.f, mx::float32, s2);
+  auto b = mx::arange(1.f, 10.f, 1.f, mx::float32, s3);
+  auto x = mx::add(a, a, s2);
+  auto y = mx::add(b, b, s3);
+
+  // The multiply will happen on the default stream.
+  std::cout << mx::multiply(x, y) << std::endl;
+
+  mx::metal::stop_capture();
+}

data/mlx/examples/cpp/timer.h

@@ -0,0 +1,20 @@
+// Copyright © 2023 Apple Inc.
+
+#pragma once
+
+#include <chrono>
+
+namespace timer {
+
+using namespace std::chrono;
+
+template <typename R, typename P>
+inline double seconds(duration<R, P> x) {
+  return duration_cast<nanoseconds>(x).count() / 1e9;
+}
+
+inline auto time() {
+  return high_resolution_clock::now();
+}
+
+} // namespace timer

data/mlx/examples/cpp/tutorial.cpp

@@ -0,0 +1,99 @@
+// Copyright © 2023 Apple Inc.
+
+#include <cassert>
+#include <iostream>
+
+#include "mlx/mlx.h"
+
+namespace mx = mlx::core;
+
+void array_basics() {
+  // Make a scalar array:
+  mx::array x(1.0);
+
+  // Get the value out of it:
+  auto s = x.item<float>();
+  assert(s == 1.0);
+
+  // Scalars have a size of 1:
+  size_t size = x.size();
+  assert(size == 1);
+
+  // Scalars have 0 dimensions:
+  int ndim = x.ndim();
+  assert(ndim == 0);
+
+  // The shape should be an empty vector:
+  auto shape = x.shape();
+  assert(shape.empty());
+
+  // The datatype should be float32:
+  auto dtype = x.dtype();
+  assert(dtype == mx::float32);
+
+  // Specify the dtype when constructing the array:
+  x = mx::array(1, mx::int32);
+  assert(x.dtype() == mx::int32);
+  x.item<int>(); // OK
+  // x.item<float>();  // Undefined!
+
+  // Make a multidimensional array:
+  x = mx::array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2});
+  // mlx is row-major by default so the first row of this array
+  // is [1.0, 2.0] and the second row is [3.0, 4.0]
+
+  // Make an array of shape {2, 2} filled with ones:
+  auto y = mx::ones({2, 2});
+
+  // Pointwise add x and y:
+  auto z = mx::add(x, y);
+
+  // Same thing:
+  z = x + y;
+
+  // mlx is lazy by default. At this point `z` only
+  // has a shape and a type but no actual data:
+  assert(z.dtype() == mx::float32);
+  assert(z.shape(0) == 2);
+  assert(z.shape(1) == 2);
+
+  // To actually run the computation you must evaluate `z`.
+  // Under the hood, mlx records operations in a graph.
+  // The variable `z` is a node in the graph which points to its operation
+  // and inputs. When `eval` is called on an array (or arrays), the array and
+  // all of its dependencies are recursively evaluated to produce the result.
+  // Once an array is evaluated, it has data and is detached from its inputs.
+  mx::eval(z);
+
+  // Of course the array can still be an input to other operations. You can
+  // even call eval on the array again, this will just be a no-op:
+  mx::eval(z); // no-op
+
+  // Some functions or methods on arrays implicitly evaluate them. For example
+  // accessing a value in an array or printing the array implicitly evaluate it:
+  z = mx::ones({1});
+  z.item<float>(); // implicit evaluation
+
+  z = mx::ones({2, 2});
+  std::cout << z << std::endl; // implicit evaluation
+}
+
+void automatic_differentiation() {
+  auto fn = [](mx::array x) { return mx::square(x); };
+
+  // Computing the derivative function of a function
+  auto grad_fn = mx::grad(fn);
+  // Call grad_fn on the input to get the derivative
+  auto x = mx::array(1.5);
+  auto dfdx = grad_fn(x);
+  // dfdx is 2 * x
+
+  // Get the second derivative by composing grad with grad
+  auto d2fdx2 = mx::grad(mx::grad(fn))(x);
+  // d2fdx2 is 2
+}
+
+int main() {
+  array_basics();
+  automatic_differentiation();
+}

data/mlx/examples/export/CMakeLists.txt

@@ -0,0 +1,22 @@
+cmake_minimum_required(VERSION 3.27)
+
+project(import_mlx LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+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)
+find_package(MLX CONFIG REQUIRED)
+
+add_executable(eval_mlp eval_mlp.cpp)
+target_link_libraries(eval_mlp PRIVATE mlx)
+
+add_executable(train_mlp train_mlp.cpp)
+target_link_libraries(train_mlp PRIVATE mlx)

data/mlx/examples/export/README.md

@@ -0,0 +1,49 @@
+## Setup
+
+Install MLX:
+
+```bash
+pip install mlx>=0.22
+```
+
+Build the C++ examples:
+
+```bash
+cmake -B build -DCMAKE_BUILD_TYPE=Release
+cmake --build build
+```
+
+## Run
+
+### Eval MLP
+
+Run the Python script to export the eval function:
+
+```bash
+python eval_mlp.py
+```
+
+Then run the C++ program to import and run the function:
+
+```
+./build/eval_mlp
+```
+
+The Python and C++ programs should output the same result.
+
+### Train MLP
+
+Run the Python script to export the model initialization and training
+functions:
+
+```bash
+python train_mlp.py
+```
+
+Then run the C++ program to import and run the functions:
+
+```
+./build/train_mlp
+```
+
+The Python and C++ programs should output the same results.

data/mlx/examples/export/eval_mlp.cpp

@@ -0,0 +1,25 @@
+// Copyright © 2024 Apple Inc.
+
+#include <mlx/mlx.h>
+#include <iostream>
+
+namespace mx = mlx::core;
+
+int main() {
+  int batch_size = 8;
+  int input_dim = 32;
+
+  // Make the input
+  mx::random::seed(42);
+  auto example_x = mx::random::uniform({batch_size, input_dim});
+
+  // Import the function
+  auto forward = mx::import_function("eval_mlp.mlxfn");
+
+  // Call the imported function
+  auto out = forward({example_x})[0];
+
+  std::cout << out << std::endl;
+
+  return 0;
+}

data/mlx/examples/export/eval_mlp.py

@@ -0,0 +1,52 @@
+# Copyright © 2024 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+
+
+class MLP(nn.Module):
+    """A simple MLP."""
+
+    def __init__(
+        self, num_layers: int, input_dim: int, hidden_dim: int, output_dim: int
+    ):
+        super().__init__()
+        layer_sizes = [input_dim] + [hidden_dim] * num_layers + [output_dim]
+        self.layers = [
+            nn.Linear(idim, odim)
+            for idim, odim in zip(layer_sizes[:-1], layer_sizes[1:])
+        ]
+
+    def __call__(self, x):
+        for l in self.layers[:-1]:
+            x = nn.relu(l(x))
+        return self.layers[-1](x)
+
+
+if __name__ == "__main__":
+
+    batch_size = 8
+    input_dim = 32
+    output_dim = 10
+
+    # Load the model
+    mx.random.seed(0)  # Seed for params
+    model = MLP(num_layers=5, input_dim=input_dim, hidden_dim=64, output_dim=output_dim)
+    mx.eval(model)
+
+    # Note, the model parameters are saved in the export function
+    def forward(x):
+        return model(x)
+
+    mx.random.seed(42)  # Seed for input
+    example_x = mx.random.uniform(shape=(batch_size, input_dim))
+
+    mx.export_function("eval_mlp.mlxfn", forward, example_x)
+
+    # Import in Python
+    imported_forward = mx.import_function("eval_mlp.mlxfn")
+    expected = forward(example_x)
+    (out,) = imported_forward(example_x)
+    assert mx.allclose(expected, out)
+    print(out)

data/mlx/examples/export/train_mlp.cpp

@@ -0,0 +1,35 @@
+// Copyright © 2024 Apple Inc.
+
+#include <mlx/mlx.h>
+#include <iostream>
+
+namespace mx = mlx::core;
+
+int main() {
+  int batch_size = 8;
+  int input_dim = 32;
+  int output_dim = 10;
+
+  auto state = mx::import_function("init_mlp.mlxfn")({});
+
+  // Make the input
+  mx::random::seed(42);
+  auto example_X = mx::random::normal({batch_size, input_dim});
+  auto example_y = mx::random::randint(0, output_dim, {batch_size});
+
+  // Import the function
+  auto step = mx::import_function("train_mlp.mlxfn");
+
+  // Call the imported function
+  for (int it = 0; it < 100; ++it) {
+    state.insert(state.end(), {example_X, example_y});
+    state = step(state);
+    eval(state);
+    auto loss = state.back();
+    state.pop_back();
+    if (it % 10 == 0) {
+      std::cout << "Loss " << loss.item<float>() << std::endl;
+    }
+  }
+  return 0;
+}