Add Metal extension build support

.github/actions/test-macos/action.yml

@@ -52,12 +52,20 @@ runs:
 
           if ${{ inputs.toolkit != 'cpu' }} ; then
             echo "::group::Build example extension"
-            pushd examples/extensions
+            pushd examples/cmake_extension
             uv pip install -r requirements.txt
             uv run setup.py build_ext --inplace
             uv run test.py
             popd
             echo "::endgroup::"
+
+            echo "::group::Build Metal extension example"
+            pushd examples/metal_extension
+            uv pip install -r requirements.txt
+            uv run --no-project --python python setup.py build_ext --inplace
+            uv run --no-project --python python test.py
+            popd
+            echo "::endgroup::"
           fi
         }
 

.gitignore

@@ -12,6 +12,7 @@ __pycache__/
 
 # Distribution / packaging
 /build/
+/examples/**/build/
 python/mlx/core
 python/mlx/share
 python/mlx/include

cmake/extension.cmake

@@ -1,47 +1,119 @@
 include(CMakeParseArguments)
 
+function(_mlx_metal_supports_logging OUTPUT_VARIABLE)
+  execute_process(
+    COMMAND xcrun -sdk macosx metal --help
+    RESULT_VARIABLE _MLX_METAL_HELP_RESULT
+    OUTPUT_VARIABLE _MLX_METAL_HELP_OUTPUT
+    ERROR_VARIABLE _MLX_METAL_HELP_ERROR
+    TIMEOUT 10)
+  set(_MLX_METAL_HELP "${_MLX_METAL_HELP_OUTPUT}\n${_MLX_METAL_HELP_ERROR}")
+  string(FIND "${_MLX_METAL_HELP}" "-fmetal-enable-logging"
+              _MLX_METAL_LOGGING_OPTION_INDEX)
+  if(_MLX_METAL_HELP_RESULT STREQUAL "0"
+     AND NOT _MLX_METAL_LOGGING_OPTION_INDEX EQUAL -1)
+    set(${OUTPUT_VARIABLE}
+        TRUE
+        PARENT_SCOPE)
+  else()
+    set(${OUTPUT_VARIABLE}
+        FALSE
+        PARENT_SCOPE)
+  endif()
+endfunction()
+
 # clang format off
 #
 # ##############################################################################
 # Build metal library
 #
-# Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib
-# from list ${SOURCES}, including list ${INCLUDE_DIRS}, depends on list ${DEPS}
+# Adds a custom target ${TARGET} to build ${OUTPUT_DIRECTORY}/{TITLE}.metallib.
+# Each source in ${SOURCES} is compiled to an AIR file before the AIR files are
+# linked into the metallib, including list ${INCLUDE_DIRS}, depends on list
+# ${DEPS}, and passes ${COMPILE_OPTIONS} to the Metal compiler.
 #
 # Args: TARGET: Custom target to be added for the metal library TITLE: Name of
 # the .metallib OUTPUT_DIRECTORY: Where to place ${TITLE}.metallib SOURCES: List
 # of source files INCLUDE_DIRS: List of include dirs DEPS: List of dependency
-# files (like headers) DEBUG: Boolean, if true, enables debug compile options
-# for this specific library. If not provided, uses global MLX_METAL_DEBUG.
+# files (like headers) COMPILE_OPTIONS: Additional Metal compiler options
+# DEPLOYMENT_TARGET: Optional macOS deployment target for this library DEBUG:
+# Boolean, if true, enables debug compile options for this specific library. If
+# not provided, uses global MLX_METAL_DEBUG.
 #
 # clang format on
 
 macro(mlx_build_metallib)
   # Parse args
-  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG)
-  set(multiValueArgs SOURCES INCLUDE_DIRS DEPS)
+  set(oneValueArgs TARGET TITLE OUTPUT_DIRECTORY DEBUG DEPLOYMENT_TARGET)
+  set(multiValueArgs SOURCES INCLUDE_DIRS DEPS COMPILE_OPTIONS)
   cmake_parse_arguments(MTLLIB "" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
 
   # Set output
   set(MTLLIB_BUILD_TARGET "${MTLLIB_OUTPUT_DIRECTORY}/${MTLLIB_TITLE}.metallib")
 
   # Collect compile options
-  set(MTLLIB_COMPILE_OPTIONS -Wall -Wextra -fno-fast-math -Wno-c++17-extensions
-                             -Wmetal-addr-spaces)
-  if(MLX_METAL_DEBUG OR MTLLIB_DEBUG)
-    set(MTLLIB_COMPILE_OPTIONS ${MTLLIB_COMPILE_OPTIONS} -gline-tables-only
-                               -frecord-sources)
+  set(_MTLLIB_COMPILE_OPTIONS
+      -Wall -Wextra -fno-fast-math -Wno-c++17-extensions -Wno-c++20-extensions
+      -Wmetal-addr-spaces)
+  if(MLX_METAL_DEBUG
+     OR MTLLIB_DEBUG
+     OR CMAKE_BUILD_TYPE STREQUAL "Debug")
+    set(_MTLLIB_COMPILE_OPTIONS ${_MTLLIB_COMPILE_OPTIONS} -gline-tables-only
+                                -frecord-sources)
+  endif()
+  if(CMAKE_BUILD_TYPE STREQUAL "Debug")
+    _mlx_metal_supports_logging(_MTLLIB_SUPPORTS_LOGGING)
+    if(_MTLLIB_SUPPORTS_LOGGING)
+      list(APPEND _MTLLIB_COMPILE_OPTIONS -fmetal-enable-logging)
+    endif()
   endif()
 
-  # Prepare metallib build command
+  set(_MTLLIB_LINK_OPTIONS)
+  if(MTLLIB_DEPLOYMENT_TARGET)
+    set(_MTLLIB_DEPLOYMENT_OPTION
+        "-mmacosx-version-min=${MTLLIB_DEPLOYMENT_TARGET}")
+    list(APPEND _MTLLIB_COMPILE_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION})
+    list(APPEND _MTLLIB_LINK_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION})
+  elseif(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "")
+    set(_MTLLIB_DEPLOYMENT_OPTION
+        "-mmacosx-version-min=${CMAKE_OSX_DEPLOYMENT_TARGET}")
+    list(APPEND _MTLLIB_COMPILE_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION})
+    list(APPEND _MTLLIB_LINK_OPTIONS ${_MTLLIB_DEPLOYMENT_OPTION})
+  endif()
+
+  # Compile each Metal source separately for incremental builds
+  set(_MTLLIB_AIR_TARGETS)
+  set(_MTLLIB_INCLUDE_OPTIONS ${MTLLIB_INCLUDE_DIRS})
+  list(TRANSFORM _MTLLIB_INCLUDE_OPTIONS PREPEND "-I")
+  set(_MTLLIB_SOURCE_INDEX 0)
+  foreach(_MTLLIB_SOURCE IN LISTS MTLLIB_SOURCES)
+    get_filename_component(_MTLLIB_SOURCE_STEM "${_MTLLIB_SOURCE}" NAME_WE)
+    set(_MTLLIB_AIR_NAME
+        "${MTLLIB_TARGET}_${_MTLLIB_SOURCE_INDEX}_${_MTLLIB_SOURCE_STEM}.air")
+    set(_MTLLIB_AIR_TARGET "${CMAKE_CURRENT_BINARY_DIR}/${_MTLLIB_AIR_NAME}")
+    set(_MTLLIB_DEPFILE "${_MTLLIB_AIR_TARGET}.d")
+    add_custom_command(
+      OUTPUT ${_MTLLIB_AIR_TARGET}
+      COMMAND
+        xcrun -sdk macosx metal ${_MTLLIB_INCLUDE_OPTIONS}
+        ${_MTLLIB_COMPILE_OPTIONS} ${MTLLIB_COMPILE_OPTIONS} -MMD -MF
+        ${_MTLLIB_DEPFILE} -MT ${_MTLLIB_AIR_TARGET} -c ${_MTLLIB_SOURCE} -o
+        ${_MTLLIB_AIR_TARGET}
+      DEPENDS ${MTLLIB_DEPS} ${_MTLLIB_SOURCE}
+      DEPFILE ${_MTLLIB_DEPFILE}
+      COMMAND_EXPAND_LISTS
+      COMMENT "Building ${_MTLLIB_AIR_NAME}"
+      VERBATIM)
+    list(APPEND _MTLLIB_AIR_TARGETS ${_MTLLIB_AIR_TARGET})
+    math(EXPR _MTLLIB_SOURCE_INDEX "${_MTLLIB_SOURCE_INDEX} + 1")
+  endforeach()
+
+  # Link the AIR files into a metallib
   add_custom_command(
     OUTPUT ${MTLLIB_BUILD_TARGET}
-    COMMAND
-      xcrun -sdk macosx metal
-      "$<LIST:TRANSFORM,${MTLLIB_INCLUDE_DIRS},PREPEND,-I>"
-      ${MTLLIB_COMPILE_OPTIONS} ${MTLLIB_SOURCES} -o ${MTLLIB_BUILD_TARGET}
-    DEPENDS ${MTLLIB_DEPS} ${MTLLIB_SOURCES}
-    COMMAND_EXPAND_LISTS
+    COMMAND xcrun -sdk macosx metal ${_MTLLIB_LINK_OPTIONS}
+            ${_MTLLIB_AIR_TARGETS} -o ${MTLLIB_BUILD_TARGET}
+    DEPENDS ${_MTLLIB_AIR_TARGETS}
     COMMENT "Building ${MTLLIB_TITLE}.metallib"
     VERBATIM)
 

docs/src/dev/extensions.rst

@@ -549,7 +549,7 @@ Building and Binding
 
 Let's look at the overall directory structure first.
 
-| extensions
+| cmake_extension
 | ├── axpby
 | │   ├── axpby.cpp
 | │   ├── axpby.h
@@ -560,13 +560,13 @@ Let's look at the overall directory structure first.
 | ├── CMakeLists.txt
 | └── setup.py
 
-* ``extensions/axpby/`` defines the C++ extension library
-* ``extensions/mlx_sample_extensions`` sets out the structure for the
+* ``cmake_extension/axpby/`` defines the C++ extension library
+* ``cmake_extension/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
+* ``cmake_extension/bindings.cpp`` provides Python bindings for our operation
+* ``cmake_extension/CMakeLists.txt`` holds CMake rules to build the library and
   Python bindings
-* ``extensions/setup.py`` holds the ``setuptools`` rules to build and install
+* ``cmake_extension/setup.py`` holds the ``setuptools`` rules to build and install
   the Python package
 
 Binding to Python
@@ -611,13 +611,6 @@ already provided, adding our :meth:`axpby` is simple.
 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
@@ -685,13 +678,12 @@ Finally, we build the nanobind_ bindings
     )
     target_link_libraries(_ext PRIVATE mlx_ext)
 
-    if(BUILD_SHARED_LIBS)
-      target_link_options(_ext PRIVATE -Wl,-rpath,@loader_path)
-    endif()
-
 Building with ``setuptools``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
+Building with a ``CMakeLists.txt``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
 Once we have set out the CMake build rules as described above, we can use the
 build utilities defined in :mod:`mlx.extension`:
 
@@ -714,21 +706,25 @@ build utilities defined in :mod:`mlx.extension`:
             python_requires=">=3.10",
         )
 
-.. note::
-    We treat ``extensions/mlx_sample_extensions`` as the package directory
-    even though it only contains a ``__init__.py`` to ensure the following:
+On macOS and Linux, ``CMakeBuild`` configures a relocatable runtime search path
+for all targets. It locates colocated libraries through ``@loader_path`` on
+macOS or ``$ORIGIN`` on Linux, and locates the MLX library under ``mlx/lib`` in
+the same ``site-packages`` directory. The extension module can therefore load
+without importing ``mlx.core`` first. Values provided through ``CMAKE_ARGS``
+can override this default.
 
-    * :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
+.. note::
+    We treat ``cmake_extension/mlx_sample_extensions`` as the package directory
+    so the C++ extension library and the Metal library are colocated 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/``)
+``python setup.py build_ext -j8 --inplace`` (in ``cmake_extension/``)
 
 This results in the directory structure:
 
-| extensions
+| cmake_extension
 | ├── mlx_sample_extensions
 | │   ├── __init__.py
 | │   ├── libmlx_ext.dylib # C++ extension library
@@ -737,11 +733,105 @@ This results in the directory structure:
 | ...
 
 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
+``cmake_extension/``), the package will be installed with the same structure as
+``cmake_extension/mlx_sample_extensions`` and the C++ and Metal library will be
 copied along with the Python binding since they are specified as
 ``package_data``.
 
+Building without a ``CMakeLists.txt``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+For a source-list based build, use ``mlx.extension.MetalExtension`` with
+``mlx.extension.BuildExtension``. The helper builds the C++/nanobind module and
+a colocated Metal library without requiring a ``CMakeLists.txt`` in the
+extension project.
+
+.. code-block:: python
+
+    from importlib.metadata import version as package_version
+
+    from mlx import extension
+    from setuptools import setup
+
+    setup(
+        name="mlx_sample_metal_extension",
+        ext_modules=[
+            extension.MetalExtension(
+                "mlx_sample_metal_extension._ext",
+                sources=[
+                    "activation.cpp",
+                    "silu_and_mul.metal",
+                    "gelu.metal",
+                    "fatrelu_and_mul.metal",
+                ],
+                extra_compile_args={
+                    "cxx": ["-O3"],
+                    "metal": ["-O3"],
+                },
+            )
+        ],
+        cmdclass={"build_ext": extension.BuildExtension},
+        install_requires=[f"mlx=={package_version('mlx')}"],
+        packages=["mlx_sample_metal_extension"],
+        package_data={"mlx_sample_metal_extension": ["*.metallib", "*.pyi"]},
+    )
+
+The final component of the extension name is also the Metal library name. The
+example above produces ``_ext.metallib``. A list passed as
+``extra_compile_args`` applies to C++; a dictionary can provide separate
+``cxx`` and ``metal`` options. The build also uses nanobind to generate
+``_ext.pyi`` next to the extension module. Include ``*.pyi`` in
+``package_data`` so that the stub is installed with the package. To disable
+stub generation, configure the command class with
+``BuildExtension.with_options(generate_stubs=False)``.
+
+Pin the ``mlx`` install requirement to the version used for the build because
+MLX does not provide a stable C++ ABI across releases. Reading the version from
+the installed package metadata avoids importing ``mlx.core`` during setup.
+
+The generated extension uses a relocatable runtime search path to locate
+``libmlx.dylib`` under ``mlx/lib`` in the same ``site-packages`` directory. This
+avoids embedding the build machine's absolute MLX library path and lets the
+extension module load without importing ``mlx.core`` first.
+
+The build defines ``MLX_EXTENSION_NAME`` as the final component of the
+extension name and ``MLX_METAL_LIBRARY_NAME`` as its string form. Use them so
+the binding does not hard-code either the module or Metal library name:
+
+.. code-block:: C++
+
+    auto library =
+        device.get_library(MLX_METAL_LIBRARY_NAME, current_binary_dir());
+
+    NB_MODULE(MLX_EXTENSION_NAME, m) {
+      // bindings
+    }
+
+Each Metal source is compiled independently to an AIR file before the AIR files
+are linked into the final metallib. Changing one Metal source therefore only
+recompiles its AIR file before relinking the metallib. Included headers are
+tracked automatically and rebuild only the AIR files that depend on them.
+
+Debug builds include Metal line tables and recorded sources. Metal logging is
+enabled when the current ``xcrun metal`` compiler advertises support. The
+effective Python
+``MACOSX_DEPLOYMENT_TARGET`` is forwarded to CMake and applied to both AIR
+compilation and metallib linking. By default, the target recorded in the
+installed ``libmlx.dylib`` is used. Set ``MACOSX_DEPLOYMENT_TARGET`` or pass
+``CMAKE_OSX_DEPLOYMENT_TARGET`` through ``CMAKE_ARGS`` to override it. This can
+select a lower target for cross-building, provided the extension is used with
+an MLX build for that target at runtime. Explicit targets below macOS 14 are
+rejected.
+
+``BuildExtension`` uses CMake's Ninja generator by default, so projects should
+include ``ninja`` in their build requirements. To use CMake's platform default
+generator instead, configure the command class with
+``BuildExtension.with_options(use_ninja=False)``. If Ninja is unavailable, the
+build emits a warning and falls back to the platform default generator.
+
+The complete standalone example is available in
+``examples/metal_extension/``.
+
 Usage
 -----
 
@@ -827,7 +917,7 @@ Scripts
 
 .. admonition:: Download the code
 
-   The full example code is available in `mlx <https://github.com/ml-explore/mlx/tree/main/examples/extensions/>`_.
+   The full example code is available in `mlx <https://github.com/ml-explore/mlx/tree/main/examples/cmake_extension/>`_.
 
 .. _Accelerate: https://developer.apple.com/documentation/accelerate/blas?language=objc
 .. _Metal: https://developer.apple.com/documentation/metal?language=objc

examples/extensions/CMakeLists.txt → examples/cmake_extension/CMakeLists.txt

@@ -80,7 +80,3 @@ nanobind_add_module(
   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()

examples/extensions/mlx_sample_extensions/__init__.py → examples/cmake_extension/mlx_sample_extensions/__init__.py

@@ -1,5 +1,3 @@
 # Copyright © 2023 Apple Inc.
 
-import mlx.core as mx
-
 from ._ext import axpby