diff --git a/.github/actions/test-wheel/action.yml b/.github/actions/test-wheel/action.yml index b745a4d2f9..35e8c3c7d5 100644 --- a/.github/actions/test-wheel/action.yml +++ b/.github/actions/test-wheel/action.yml @@ -10,15 +10,24 @@ inputs: runs: using: 'composite' steps: - - name: Get Python version + - name: Get frontend ABI id: python shell: bash - run: python -c "import sys; print(f'version={sys.version_info.major}.{sys.version_info.minor}')" >> $GITHUB_OUTPUT + run: | + python - <<'PY' >> "$GITHUB_OUTPUT" + import sys + import sysconfig + + if sysconfig.get_config_var("Py_GIL_DISABLED"): + print(f"abi=py{sys.version_info.major}.{sys.version_info.minor}t") + else: + print("abi=abi3") + PY - name: Download frontend packages uses: actions/download-artifact@v8 with: - pattern: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ steps.python.outputs.version }} + pattern: frontend-${{ runner.os }}-${{ runner.arch }}-${{ steps.python.outputs.abi }} path: wheelhouse - name: Download backend packages diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index b53f003ff4..2b2aa91380 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -58,12 +58,18 @@ jobs: matrix: os: ['Linux', 'Windows'] arch: ['x86_64', 'aarch64'] - python-version: &pyver ['3.10', '3.11', '3.12', '3.13', '3.13t', '3.14', '3.14t'] - # There is no cp310 binary for Windows on arm. + # One Stable ABI frontend wheel covers all supported regular CPython + # versions on a platform. Python 3.13t and 3.14t do not have a + # free-threaded Stable ABI and require version-specific linked wheels. + python-version: &pyver ['3.10', '3.13t', '3.14t'] exclude: - os: 'Windows' arch: 'aarch64' python-version: '3.10' + include: + - os: 'Windows' + arch: 'aarch64' + python-version: '3.11' runs-on: |- ${{ case(matrix.os == 'Windows', case(matrix.arch == 'aarch64', 'windows-11-arm', @@ -89,7 +95,7 @@ jobs: - run: unzip -l wheelhouse/*.whl - uses: actions/upload-artifact@v7 with: - name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} + name: frontend-${{ runner.os }}-${{ runner.arch }}-${{ contains(matrix.python-version, 't') && format('py{0}', matrix.python-version) || 'abi3' }} path: wheelhouse/mlx-*.whl if-no-files-found: error @@ -175,7 +181,7 @@ jobs: - name: Upload frontend packages uses: actions/upload-artifact@v7 with: - name: frontend-${{ runner.os }}-${{ runner.arch }}-py${{ matrix.python-version }} + name: frontend-${{ runner.os }}-${{ runner.arch }}-${{ contains(matrix.python-version, 't') && format('py{0}', matrix.python-version) || 'abi3' }} path: wheelhouse/mlx-*.whl if-no-files-found: error - name: Upload backend packages @@ -187,30 +193,36 @@ jobs: if-no-files-found: error test_wheel: - name: Test (${{ matrix.os }}, ${{ matrix.toolkit }}, ${{ matrix.arch }}) + name: Test (${{ matrix.os }}, python-${{ matrix.python-version }}, ${{ matrix.toolkit }}, ${{ matrix.arch }}) if: github.repository == 'ml-explore/mlx' needs: [build_frontend, build_backend, build_mac_wheels] strategy: matrix: os: ['Linux', 'Windows'] - arch: ['aarch64'] + arch: ['x86_64', 'aarch64'] toolkit: ['cpu'] + python-version: ['3.12', '3.13t', '3.14t'] include: - - os: 'Linux' - arch: 'x86_64' - toolkit: 'cpu' - os: 'Linux' arch: 'x86_64' toolkit: 'cuda-12.9' + python-version: '3.12' - os: 'Linux' arch: 'x86_64' toolkit: 'cuda-13.0' - - os: 'Windows' - arch: 'x86_64' - toolkit: 'cpu' + python-version: '3.12' - os: 'macOS' arch: 'aarch64' toolkit: 'metal' + python-version: '3.12' + - os: 'macOS' + arch: 'aarch64' + toolkit: 'metal' + python-version: '3.13t' + - os: 'macOS' + arch: 'aarch64' + toolkit: 'metal' + python-version: '3.14t' runs-on: |- ${{ case(matrix.os == 'Windows', case(matrix.arch == 'aarch64', 'windows-11-arm', 'windows-2022'), @@ -224,6 +236,7 @@ jobs: - uses: ./.github/actions/setup with: toolkit: ${{ matrix.toolkit }} + python-version: ${{ matrix.python-version }} use-ccache: false - uses: ./.github/actions/test-wheel with: diff --git a/CMakeLists.txt b/CMakeLists.txt index feb7ce5ebb..07d6e00d8d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,4 +1,4 @@ -cmake_minimum_required(VERSION 3.25) +cmake_minimum_required(VERSION 3.26) if(NOT MLX_VERSION) file(STRINGS "mlx/version.h" _mlx_h_version REGEX "^#define MLX_VERSION_.*$") @@ -392,10 +392,26 @@ if(MLX_BUILD_PYTHON_BINDINGS) Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) + execute_process( + COMMAND + "${Python_EXECUTABLE}" -c + "import sysconfig; print('ON' if sysconfig.get_config_var('Py_GIL_DISABLED') else 'OFF')" + RESULT_VARIABLE _mlx_python_gil_result + OUTPUT_VARIABLE MLX_PYTHON_FREE_THREADED + OUTPUT_STRIP_TRAILING_WHITESPACE) + if(NOT _mlx_python_gil_result EQUAL 0) + message(FATAL_ERROR "Failed to determine whether Python is free-threaded.") + endif() + if(NOT MLX_PYTHON_FREE_THREADED) + find_package( + Python 3.10 + COMPONENTS Development.SABIModule + REQUIRED) + endif() FetchContent_Declare( nanobind GIT_REPOSITORY https://github.com/wjakob/nanobind.git - GIT_TAG v2.15.0 + GIT_TAG v3.0.0 GIT_SHALLOW TRUE EXCLUDE_FROM_ALL) FetchContent_MakeAvailable(nanobind) diff --git a/MANIFEST.in b/MANIFEST.in index 632fae78c6..6a3cf1af8a 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,5 +1,6 @@ include CMakeLists.txt include mlx.pc.in +include _build_backend/backend.py recursive-include mlx * include cmake/* include python/src/* diff --git a/_build_backend/backend.py b/_build_backend/backend.py new file mode 100644 index 0000000000..751f612780 --- /dev/null +++ b/_build_backend/backend.py @@ -0,0 +1,23 @@ +import os +import sysconfig + +from setuptools import build_meta as _setuptools_build_meta +from setuptools.build_meta import * # type: ignore # noqa: F403 + + +def _with_nanobind_backend(requires): + build_backend = int(os.environ.get("MLX_BUILD_BACKEND_PACKAGE", 0)) + free_threaded = bool(sysconfig.get_config_var("Py_GIL_DISABLED")) + if not build_backend and not free_threaded: + requires.append("nanobind-backend>=1.0.0") + return requires + + +def get_requires_for_build_wheel(config_settings=None): + requires = _setuptools_build_meta.get_requires_for_build_wheel(config_settings) + return _with_nanobind_backend(requires) + + +def get_requires_for_build_editable(config_settings=None): + requires = _setuptools_build_meta.get_requires_for_build_editable(config_settings) + return _with_nanobind_backend(requires) diff --git a/docs/src/dev/extensions.rst b/docs/src/dev/extensions.rst index 89c8991e56..4722e886a7 100644 --- a/docs/src/dev/extensions.rst +++ b/docs/src/dev/extensions.rst @@ -679,7 +679,8 @@ Finally, we build the nanobind_ bindings nanobind_add_module( _ext - NB_STATIC STABLE_ABI LTO NOMINSIZE + BACKEND_MODULE nanobind_backend + LTO NOMINSIZE NB_DOMAIN mlx ${CMAKE_CURRENT_LIST_DIR}/bindings.cpp ) @@ -709,9 +710,11 @@ build utilities defined in :mod:`mlx.extension`: cmdclass={"build_ext": extension.CMakeBuild}, packages=["mlx_sample_extensions"], package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, + install_requires=["nanobind-backend>=1.0.0"], + options={"bdist_wheel": {"py_limited_api": "cp310"}}, extras_require={"dev":[]}, zip_safe=False, - python_requires=">=3.8", + python_requires=">=3.10", ) .. note:: @@ -733,7 +736,7 @@ This results in the directory structure: | │ ├── __init__.py | │ ├── libmlx_ext.dylib # C++ extension library | │ ├── mlx_ext.metallib # Metal library -| │ └── _ext.cpython-3x-darwin.so # Python Binding +| │ └── _ext.abi3.so # Python Binding | ... When you try to install using the command ``python -m pip install .`` (in diff --git a/examples/extensions/CMakeLists.txt b/examples/extensions/CMakeLists.txt index eedc3a7dab..d7073b782d 100644 --- a/examples/extensions/CMakeLists.txt +++ b/examples/extensions/CMakeLists.txt @@ -11,8 +11,8 @@ option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON) # ----------------------------- Dependencies ----------------------------- find_package( - Python 3.8 - COMPONENTS Interpreter Development.Module + Python 3.10 + COMPONENTS Interpreter Development.Module Development.SABIModule REQUIRED) execute_process( COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir @@ -64,8 +64,8 @@ endif() # ----------------------------- Python Bindings ----------------------------- nanobind_add_module( _ext - NB_STATIC - STABLE_ABI + BACKEND_MODULE + nanobind_backend LTO NOMINSIZE NB_DOMAIN diff --git a/examples/extensions/pyproject.toml b/examples/extensions/pyproject.toml index 560a58bc28..04185bc5e7 100644 --- a/examples/extensions/pyproject.toml +++ b/examples/extensions/pyproject.toml @@ -1,8 +1,9 @@ [build-system] requires = [ "setuptools>=42", - "cmake>=3.25", + "cmake>=3.27", "mlx>=0.18.0", - "nanobind==2.15.0", + "nanobind==3.0.0", + "nanobind-backend>=1.0.0", ] build-backend = "setuptools.build_meta" diff --git a/examples/extensions/requirements.txt b/examples/extensions/requirements.txt index 917d125eea..635b14ac28 100644 --- a/examples/extensions/requirements.txt +++ b/examples/extensions/requirements.txt @@ -1,4 +1,5 @@ setuptools>=42 -cmake>=3.25 +cmake>=3.27 mlx>=0.31.2 -nanobind==2.15.0 +nanobind==3.0.0 +nanobind-backend>=1.0.0 diff --git a/examples/extensions/setup.py b/examples/extensions/setup.py index be990a02b6..4bffae6d5c 100644 --- a/examples/extensions/setup.py +++ b/examples/extensions/setup.py @@ -13,6 +13,8 @@ cmdclass={"build_ext": extension.CMakeBuild}, packages=["mlx_sample_extensions"], package_data={"mlx_sample_extensions": ["*.so", "*.dylib", "*.metallib"]}, + install_requires=["nanobind-backend>=1.0.0"], + options={"bdist_wheel": {"py_limited_api": "cp310"}}, zip_safe=False, python_requires=">=3.10", ) diff --git a/pyproject.toml b/pyproject.toml index 45675ae856..f8f90c0e56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,8 @@ [build-system] requires = [ "setuptools>=80", - "cmake>=3.25", + "cmake>=3.26", "typing_extensions", ] -build-backend = "setuptools.build_meta" +build-backend = "backend" +backend-path = ["_build_backend"] diff --git a/python/mlx/extension.py b/python/mlx/extension.py index d794fbe28f..c352ee4b42 100644 --- a/python/mlx/extension.py +++ b/python/mlx/extension.py @@ -17,7 +17,7 @@ # A CMakeExtension needs a sourcedir instead of a file list. class CMakeExtension(Extension): def __init__(self, name: str, sourcedir: str = "") -> None: - super().__init__(name, sources=[]) + super().__init__(name, sources=[], py_limited_api=True) self.sourcedir = os.fspath(Path(sourcedir).resolve()) diff --git a/python/src/CMakeLists.txt b/python/src/CMakeLists.txt index 0798add410..141a35ec31 100644 --- a/python/src/CMakeLists.txt +++ b/python/src/CMakeLists.txt @@ -1,8 +1,14 @@ +# Python 3.13t and 3.14t have no free-threaded Stable ABI, so they must use +# linked mode instead of the regular split-mode build. +if(MLX_PYTHON_FREE_THREADED) + set(MLX_NANOBIND_MODE NB_STATIC FREE_THREADED) +else() + set(MLX_NANOBIND_MODE BACKEND_MODULE nanobind_backend) +endif() + nanobind_add_module( core - NB_STATIC - STABLE_ABI - FREE_THREADED + ${MLX_NANOBIND_MODE} LTO NOMINSIZE NB_DOMAIN diff --git a/python/src/buffer.h b/python/src/buffer.h index 4b194b3d25..48dbf4ecfc 100644 --- a/python/src/buffer.h +++ b/python/src/buffer.h @@ -7,6 +7,28 @@ #include "mlx/array.h" #include "mlx/utils.h" +// Py_buffer became part of the Stable ABI in Python 3.11. CPython 3.10 uses +// the same layout, which split-mode extensions targeting cp310 must declare. +#if defined(Py_LIMITED_API) && Py_LIMITED_API + 0 < 0x030B0000 +typedef struct { + void* buf; + PyObject* obj; + Py_ssize_t len; + Py_ssize_t itemsize; + int readonly; + int ndim; + char* format; + Py_ssize_t* shape; + Py_ssize_t* strides; + Py_ssize_t* suboffsets; + void* internal; +} Py_buffer; + +#define PyBUF_FORMAT 0x0004 +#define PyBUF_ND 0x0008 +#define PyBUF_STRIDES (0x0010 | PyBUF_ND) +#endif + // Only defined in >= Python 3.9 // https://github.com/python/cpython/blob/f6cdc6b4a191b75027de342aa8b5d344fb31313e/Include/typeslots.h#L2-L3 #ifndef Py_bf_getbuffer diff --git a/python/src/convert.h b/python/src/convert.h index 133f443ed7..7c767a5944 100644 --- a/python/src/convert.h +++ b/python/src/convert.h @@ -14,24 +14,6 @@ namespace nb = nanobind; namespace nanobind { -template <> -struct ndarray_traits { - static constexpr bool is_complex = false; - static constexpr bool is_float = true; - static constexpr bool is_bool = false; - static constexpr bool is_int = false; - static constexpr bool is_signed = true; -}; - -template <> -struct ndarray_traits { - static constexpr bool is_complex = false; - static constexpr bool is_float = true; - static constexpr bool is_bool = false; - static constexpr bool is_int = false; - static constexpr bool is_signed = true; -}; - namespace detail { template <> diff --git a/python/src/load.cpp b/python/src/load.cpp index 605049864e..39e8a1ab3a 100644 --- a/python/src/load.cpp +++ b/python/src/load.cpp @@ -15,6 +15,12 @@ #include "python/src/small_vector.h" #include "python/src/utils.h" +// These flags became part of the Stable ABI together with Py_buffer in 3.11. +#if defined(Py_LIMITED_API) && Py_LIMITED_API + 0 < 0x030B0000 +#define PyBUF_READ 0x100 +#define PyBUF_WRITE 0x200 +#endif + namespace mx = mlx::core; namespace nb = nanobind; using namespace nb::literals; diff --git a/python/src/mlx_func.cpp b/python/src/mlx_func.cpp index 9955e134f9..3bec0ccf14 100644 --- a/python/src/mlx_func.cpp +++ b/python/src/mlx_func.cpp @@ -49,7 +49,11 @@ PyObject* gc_func_vectorcall( PyObject* const* args, size_t nargs, PyObject* kwnames) { - return PyObject_Vectorcall(((gc_func*)self)->func, args, nargs, kwnames); + return nb::detail::vectorcall(((gc_func*)self)->func, args, nargs, kwnames); +} + +PyObject* gc_func_call(PyObject* self, PyObject* args, PyObject* kwargs) { + return PyObject_Call(((gc_func*)self)->func, args, kwargs); } void gc_func_dealloc(PyObject* self) { @@ -83,16 +87,22 @@ PyType_Slot gc_func_slots[] = { {Py_tp_getset, (void*)gc_func_getset}, {Py_tp_getattro, (void*)gc_func_getattro}, {Py_tp_members, (void*)gc_func_members}, - {Py_tp_call, (void*)PyVectorcall_Call}, + {Py_tp_call, (void*)gc_func_call}, {Py_tp_dealloc, (void*)gc_func_dealloc}, {0, 0}}; +#if defined(Py_TPFLAGS_HAVE_VECTORCALL) +constexpr auto gc_func_vectorcall_flag = Py_TPFLAGS_HAVE_VECTORCALL; +#else +constexpr unsigned long gc_func_vectorcall_flag = 1UL << 11; +#endif + static PyType_Spec gc_func_spec = { /* .name = */ "mlx.gc_func", /* .basicsize = */ (int)sizeof(gc_func), /* .itemsize = */ 0, /* .flags = */ Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | - Py_TPFLAGS_HAVE_VECTORCALL, + gc_func_vectorcall_flag, /* .slots = */ gc_func_slots}; static PyTypeObject* gc_func_tp = nullptr; diff --git a/python/src/small_vector.h b/python/src/small_vector.h index 0b2e6bb2ee..3744ea10f0 100644 --- a/python/src/small_vector.h +++ b/python/src/small_vector.h @@ -31,13 +31,13 @@ struct type_caster<::mlx::core::SmallVector> { // Not noexcept: on overflow of a narrow integer element we raise // OverflowError so nanobind surfaces a clean error to the user. - bool from_python(handle src, uint8_t flags, cleanup_list* cleanup) { + bool from_python(handle src, uint32_t flags, cleanup_list* cleanup) { size_t size; PyObject* temp; // Will initialize 'size' and 'temp'. All return values and // return parameters are zero/NULL in the case of a failure. - PyObject** o = seq_get(src.ptr(), &size, &temp); + PyObject** o = NB_CALL(seq_get)(src.ptr(), &size, &temp); value.clear(); value.reserve(size); diff --git a/python/src/transforms.cpp b/python/src/transforms.cpp index 1ec20a1375..2d9e2d3467 100644 --- a/python/src/transforms.cpp +++ b/python/src/transforms.cpp @@ -809,8 +809,8 @@ class PyCustomFunction { } int array_index = 0; int tangent_index = 0; - auto new_tangents = - nb::cast(tree_map(args, [&](nb::handle element) { + auto new_tangents = nb::cast( + tree_map(args, [&](nb::handle element) -> nb::object { if (nb::isinstance(element) && have_tangents[array_index++]) { return nb::cast(tangents[tangent_index++]); @@ -856,8 +856,8 @@ class PyCustomFunction { } int arr_index = 0; - auto new_axes = - nb::cast(tree_map(args, [&](nb::handle element) { + auto new_axes = nb::cast( + tree_map(args, [&](nb::handle element) -> nb::object { int axis = axes[arr_index++]; if (nb::isinstance(element) && axis >= 0) { return nb::cast(axis); diff --git a/setup.py b/setup.py index 52cc4f7492..3aab23d594 100644 --- a/setup.py +++ b/setup.py @@ -6,6 +6,7 @@ import re import subprocess import sys +import sysconfig from functools import partial from pathlib import Path @@ -71,6 +72,8 @@ def get_version(): build_backend = int(os.environ.get("MLX_BUILD_BACKEND_PACKAGE", 0)) build_macos = platform.system() == "Darwin" build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "") +# Free-threaded wheels use the interpreter-specific linked ABI. +free_threaded = bool(sysconfig.get_config_var("Py_GIL_DISABLED")) # A CMakeExtension needs a sourcedir instead of a file list. @@ -78,7 +81,7 @@ def get_version(): # If you need multiple extensions, see scikit-build. class CMakeExtension(Extension): def __init__(self, name: str, sourcedir: str = "") -> None: - super().__init__(name, sources=[]) + super().__init__(name, sources=[], py_limited_api=not free_threaded) self.sourcedir = os.fspath(Path(sourcedir).resolve()) @@ -208,6 +211,12 @@ def get_tag(self) -> tuple[str, str, str]: if build_backend: impl = self.python_tag abi = "none" + elif not free_threaded: + # nanobind split mode targets the Python 3.10 Stable ABI. The + # official backend has no CPython 3.10 wheel for Windows ARM64, so + # that platform starts at CPython 3.11 instead. + impl = "cp311" if plat_name == "win_arm64" else "cp310" + abi = "abi3" return (impl, abi, plat_name) def write_wheelfile(self, *args, **kwargs) -> None: @@ -299,6 +308,8 @@ def is_backend_file(file): install_requires = [] if not build_backend: + if not free_threaded: + install_requires.append("nanobind-backend>=1.0.0") if build_frontend: install_requires.append( f'mlx-metal=={version}; platform_system == "Darwin"'