From c9d30aa6acffa63caf142e9e0a3689723342fdb0 Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Thu, 2 Jan 2025 19:09:04 -0800 Subject: [PATCH] MLX in C++ example (#1736) * MLX in C++ example * nits * fix docs --- docs/src/dev/mlx_in_cpp.rst | 121 ++++++++++++++++++++++++++ docs/src/index.rst | 1 + docs/src/install.rst | 2 + docs/src/python/ops.rst | 1 + examples/cmake_project/CMakeLists.txt | 22 +++++ examples/cmake_project/README.md | 26 ++++++ examples/cmake_project/example.cpp | 14 +++ examples/export/CMakeLists.txt | 7 +- examples/export/eval_mlp.cpp | 8 +- examples/export/train_mlp.cpp | 12 +-- mlx/ops.h | 2 +- python/mlx/__main__.py | 27 ++++++ python/src/ops.cpp | 30 ++++--- 13 files changed, 242 insertions(+), 31 deletions(-) create mode 100644 docs/src/dev/mlx_in_cpp.rst create mode 100644 examples/cmake_project/CMakeLists.txt create mode 100644 examples/cmake_project/README.md create mode 100644 examples/cmake_project/example.cpp create mode 100644 python/mlx/__main__.py diff --git a/docs/src/dev/mlx_in_cpp.rst b/docs/src/dev/mlx_in_cpp.rst new file mode 100644 index 000000000..af778b33b --- /dev/null +++ b/docs/src/dev/mlx_in_cpp.rst @@ -0,0 +1,121 @@ +.. _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 17) + 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 diff --git a/docs/src/index.rst b/docs/src/index.rst index 4c41800d7..99cb7a8af 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -87,3 +87,4 @@ are the CPU and GPU. dev/extensions dev/metal_debugger dev/custom_metal_kernels + dev/mlx_in_cpp diff --git a/docs/src/install.rst b/docs/src/install.rst index 5b9e93216..059b2cba4 100644 --- a/docs/src/install.rst +++ b/docs/src/install.rst @@ -1,3 +1,5 @@ +.. _build_and_install: + Build and Install ================= diff --git a/docs/src/python/ops.rst b/docs/src/python/ops.rst index 61dbc7596..9e50feebb 100644 --- a/docs/src/python/ops.rst +++ b/docs/src/python/ops.rst @@ -89,6 +89,7 @@ Operations isneginf isposinf issubdtype + kron left_shift less less_equal diff --git a/examples/cmake_project/CMakeLists.txt b/examples/cmake_project/CMakeLists.txt new file mode 100644 index 000000000..7a73ce1d9 --- /dev/null +++ b/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) diff --git a/examples/cmake_project/README.md b/examples/cmake_project/README.md new file mode 100644 index 000000000..657c58fe6 --- /dev/null +++ b/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) +``` diff --git a/examples/cmake_project/example.cpp b/examples/cmake_project/example.cpp new file mode 100644 index 000000000..6e69f47c5 --- /dev/null +++ b/examples/cmake_project/example.cpp @@ -0,0 +1,14 @@ +// Copyright © 2024 Apple Inc. + +#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; +} diff --git a/examples/export/CMakeLists.txt b/examples/export/CMakeLists.txt index c30011406..aefcdf59f 100644 --- a/examples/export/CMakeLists.txt +++ b/examples/export/CMakeLists.txt @@ -2,20 +2,15 @@ cmake_minimum_required(VERSION 3.27) project(import_mlx LANGUAGES CXX) -# ----------------------------- Setup ----------------------------- set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) -set(CMAKE_POSITION_INDEPENDENT_CODE ON) -# ----------------------------- Dependencies ----------------------------- find_package( Python 3.9 COMPONENTS Interpreter Development.Module REQUIRED) execute_process( - COMMAND "${Python_EXECUTABLE}" -m pip show mlx - COMMAND grep location - COMMAND awk "{print $4 \"/mlx\"}" + COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE MLX_ROOT) find_package(MLX CONFIG REQUIRED) diff --git a/examples/export/eval_mlp.cpp b/examples/export/eval_mlp.cpp index 2facae43a..336cdd272 100644 --- a/examples/export/eval_mlp.cpp +++ b/examples/export/eval_mlp.cpp @@ -3,18 +3,18 @@ #include #include -using namespace mlx::core; +namespace mx = mlx::core; int main() { int batch_size = 8; int input_dim = 32; // Make the input - random::seed(42); - auto example_x = random::uniform({batch_size, input_dim}); + mx::random::seed(42); + auto example_x = mx::random::uniform({batch_size, input_dim}); // Import the function - auto forward = import_function("eval_mlp.mlxfn"); + auto forward = mx::import_function("eval_mlp.mlxfn"); // Call the imported function auto out = forward({example_x})[0]; diff --git a/examples/export/train_mlp.cpp b/examples/export/train_mlp.cpp index c3d516e9e..5950bebd3 100644 --- a/examples/export/train_mlp.cpp +++ b/examples/export/train_mlp.cpp @@ -3,22 +3,22 @@ #include #include -using namespace mlx::core; +namespace mx = mlx::core; int main() { int batch_size = 8; int input_dim = 32; int output_dim = 10; - auto state = import_function("init_mlp.mlxfn")({}); + auto state = mx::import_function("init_mlp.mlxfn")({}); // Make the input - random::seed(42); - auto example_X = random::normal({batch_size, input_dim}); - auto example_y = random::randint(0, output_dim, {batch_size}); + 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 = import_function("train_mlp.mlxfn"); + auto step = mx::import_function("train_mlp.mlxfn"); // Call the imported function for (int it = 0; it < 100; ++it) { diff --git a/mlx/ops.h b/mlx/ops.h index ec8cf20e2..aadd187cb 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -914,7 +914,7 @@ inline array gather( return gather(a, {indices}, std::vector{axis}, slice_sizes, s); } -/** Returns Kronecker Producct given two input arrays. */ +/** Compute the Kronecker product of two arrays. */ array kron(const array& a, const array& b, StreamOrDevice s = {}); /** Take array slices at the given indices of the specified axis. */ diff --git a/python/mlx/__main__.py b/python/mlx/__main__.py new file mode 100644 index 000000000..6c21ef112 --- /dev/null +++ b/python/mlx/__main__.py @@ -0,0 +1,27 @@ +import argparse + + +def main() -> None: + from mlx.core import __version__ + + parser = argparse.ArgumentParser() + parser.add_argument( + "--version", + action="version", + version=__version__, + help="Print the version number.", + ) + parser.add_argument( + "--cmake-dir", + action="store_true", + help="Print the path to the MLX CMake module directory.", + ) + args = parser.parse_args() + if args.cmake_dir: + from pathlib import Path + + print(Path(__file__).parent) + + +if __name__ == "__main__": + main() diff --git a/python/src/ops.cpp b/python/src/ops.cpp index b58ee670f..e0315a437 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1468,24 +1468,26 @@ void init_ops(nb::module_& m) { nb::sig( "def kron(a: array, b: array, *, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( - Compute the Kronecker product of two arrays `a` and `b`. + Compute the Kronecker product of two arrays ``a`` and ``b``. + Args: - a (array): The first input array - b (array): The second input array - stream (Union[None, Stream, Device], optional): Optional stream or device for execution. - Default is `None`. + a (array): The first input array. + b (array): The second input array. + stream (Union[None, Stream, Device], optional): Optional stream or + device for execution. Default: ``None``. + Returns: - array: The Kronecker product of `a` and `b`. + array: The Kronecker product of ``a`` and ``b``. + Examples: - >>> import mlx - >>> a = mlx.array([[1, 2], [3, 4]]) - >>> b = mlx.array([[0, 5], [6, 7]]) - >>> result = mlx.kron(a, b) + >>> a = mx.array([[1, 2], [3, 4]]) + >>> b = mx.array([[0, 5], [6, 7]]) + >>> result = mx.kron(a, b) >>> print(result) - [[ 0 5 0 10] - [ 6 7 12 14] - [ 0 15 0 20] - [18 21 24 28]] + array([[0, 5, 0, 10], + [6, 7, 12, 14], + [0, 15, 0, 20], + [18, 21, 24, 28]], dtype=int32) )pbdoc"); m.def( "take",