mlx 1.0.0

data/mlx/setup.py

@@ -0,0 +1,315 @@
+# Copyright © 2023 Apple Inc.
+
+import datetime
+import os
+import platform
+import re
+import subprocess
+from functools import partial
+from pathlib import Path
+
+from setuptools import Extension, find_namespace_packages, setup
+from setuptools.command.bdist_wheel import bdist_wheel
+from setuptools.command.build_ext import build_ext
+
+
+def cuda_toolkit_major_version():
+    out = subprocess.check_output(["nvcc", "--version"], stderr=subprocess.STDOUT)
+    text = out.decode()
+    m = re.search(r"release (\d+)", text)
+    if m:
+        return int(m.group(1))
+    return None
+
+
+def get_version():
+    with open("mlx/version.h", "r") as fid:
+        for l in fid:
+            if "#define MLX_VERSION_MAJOR" in l:
+                major = l.split()[-1]
+            if "#define MLX_VERSION_MINOR" in l:
+                minor = l.split()[-1]
+            if "#define MLX_VERSION_PATCH" in l:
+                patch = l.split()[-1]
+    version = f"{major}.{minor}.{patch}"
+    pypi_release = int(os.environ.get("PYPI_RELEASE", 0))
+    dev_release = int(os.environ.get("DEV_RELEASE", 0))
+    if not pypi_release or dev_release:
+        today = datetime.date.today()
+        version = f"{version}.dev{today.year}{today.month:02d}{today.day:02d}"
+    if not pypi_release and not dev_release:
+        git_hash = (
+            subprocess.run(
+                "git rev-parse --short HEAD".split(),
+                capture_output=True,
+                check=True,
+            )
+            .stdout.strip()
+            .decode()
+        )
+        version = f"{version}+{git_hash}"
+
+    return version
+
+
+build_stage = int(os.environ.get("MLX_BUILD_STAGE", 0))
+build_macos = platform.system() == "Darwin"
+build_cuda = "MLX_BUILD_CUDA=ON" in os.environ.get("CMAKE_ARGS", "")
+
+
+# A CMakeExtension needs a sourcedir instead of a file list.
+# The name must be the _single_ output extension from the CMake build.
+# If you need multiple extensions, see scikit-build.
+class CMakeExtension(Extension):
+    def __init__(self, name: str, sourcedir: str = "") -> None:
+        super().__init__(name, sources=[])
+        self.sourcedir = os.fspath(Path(sourcedir).resolve())
+
+
+class CMakeBuild(build_ext):
+    def build_extension(self, ext: CMakeExtension) -> None:
+        # Must be in this form due to bug in .resolve() only fixed in Python 3.10+
+        ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)  # type: ignore[no-untyped-call]
+        extdir = ext_fullpath.parent.resolve()
+
+        debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug
+        cfg = "Debug" if debug else "Release"
+
+        build_temp = Path(self.build_temp) / ext.name
+        if not build_temp.exists():
+            build_temp.mkdir(parents=True)
+
+        install_prefix = extdir
+        pybind_out_dir = extdir
+        if build_stage == 1:
+            # Don't include MLX libraries in the wheel
+            install_prefix = build_temp
+        elif build_stage == 2:
+            # Don't include Python bindings in the wheel
+            pybind_out_dir = build_temp
+        cmake_args = [
+            f"-DCMAKE_INSTALL_PREFIX={install_prefix}",
+            f"-DMLX_PYTHON_BINDINGS_OUTPUT_DIRECTORY={pybind_out_dir}",
+            f"-DCMAKE_BUILD_TYPE={cfg}",
+            "-DMLX_BUILD_PYTHON_BINDINGS=ON",
+            "-DMLX_BUILD_TESTS=OFF",
+            "-DMLX_BUILD_BENCHMARKS=OFF",
+            "-DMLX_BUILD_EXAMPLES=OFF",
+            "-DBUILD_SHARED_LIBS=ON",
+        ]
+        if build_stage == 2 and build_cuda:
+            # Last arch is always real and virtual for forward-compatibility
+            cuda_archs = ";".join(
+                (
+                    "75-real",
+                    "80-real",
+                    "90a-real",
+                    "100a-real",
+                    "120a-real",
+                    "120-virtual",
+                )
+            )
+            cmake_args += [f"-DMLX_CUDA_ARCHITECTURES={cuda_archs}"]
+            # Search CUDA libs from python packages.
+            cmake_args += ["-DMLX_LOAD_CUDA_LIBS_FROM_PYTHON=ON"]
+
+        # Some generators require explcitly passing config when building.
+        build_args = ["--config", cfg]
+        # Adding CMake arguments set as environment variable
+        # (needed e.g. to build for ARM OSx on conda-forge)
+        if "CMAKE_ARGS" in os.environ:
+            cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item]
+
+        # Pass version to C++
+        cmake_args += [f"-DMLX_VERSION={self.distribution.get_version()}"]  # type: ignore[attr-defined]
+
+        if build_macos:
+            # Cross-compile support for macOS - respect ARCHFLAGS if set
+            archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", ""))
+            if archs:
+                cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))]
+
+        # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level
+        # across all generators.
+        if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ:
+            build_args += [f"-j{os.cpu_count()}"]
+
+        # Avoid cache miss when building from temporary dirs.
+        os.environ["CCACHE_BASEDIR"] = os.path.realpath(self.build_temp)
+        os.environ["CCACHE_NOHASHDIR"] = "true"
+
+        subprocess.run(
+            ["cmake", ext.sourcedir, *cmake_args], cwd=build_temp, check=True
+        )
+        subprocess.run(
+            ["cmake", "--build", ".", "--target", "install", *build_args],
+            cwd=build_temp,
+            check=True,
+        )
+
+    # Make sure to copy mlx.metallib for inplace builds
+    def run(self):
+        super().run()
+
+        ext = next(ext for ext in self.extensions if ext.name == "mlx.core")
+
+        # Based on https://github.com/pypa/setuptools/blob/main/setuptools/command/build_ext.py#L102
+        if self.inplace:
+            # Resolve inplace package dir
+            build_py = self.get_finalized_command("build_py")
+            inplace_file, regular_file = self._get_inplace_equivalent(build_py, ext)
+
+            inplace_dir = str(Path(inplace_file).parent.resolve())
+            regular_dir = str(Path(regular_file).parent.resolve())
+
+            self.copy_tree(regular_dir, inplace_dir)
+
+        # Build type stubs.
+        build_temp = Path(self.build_temp) / ext.name
+        subprocess.run(
+            ["cmake", "--install", build_temp, "--component", "core_stub"],
+            check=True,
+        )
+
+
+class MLXBdistWheel(bdist_wheel):
+    def get_tag(self) -> tuple[str, str, str]:
+        impl, abi, plat_name = super().get_tag()
+        if build_stage == 2:
+            impl = self.python_tag
+            abi = "none"
+        return (impl, abi, plat_name)
+
+
+# Read the content of README.md
+with open(Path(__file__).parent / "README.md", encoding="utf-8") as f:
+    long_description = f.read()
+
+
+if __name__ == "__main__":
+    package_dir = {"": "python"}
+    packages = find_namespace_packages(
+        where="python",
+        exclude=[
+            "src",
+            "tests",
+            "scripts",
+            "mlx.lib",
+            "mlx.include",
+            "mlx.share",
+            "mlx.share.**",
+            "mlx.include.**",
+        ],
+    )
+
+    version = get_version()
+
+    _setup = partial(
+        setup,
+        version=version,
+        author="MLX Contributors",
+        author_email="mlx@group.apple.com",
+        description="A framework for machine learning on Apple silicon.",
+        long_description=long_description,
+        long_description_content_type="text/markdown",
+        license="MIT",
+        url="https://github.com/ml-explore/mlx",
+        include_package_data=True,
+        package_dir=package_dir,
+        zip_safe=False,
+        python_requires=">=3.10",
+        ext_modules=[CMakeExtension("mlx.core")],
+        cmdclass={
+            "build_ext": CMakeBuild,
+            "bdist_wheel": MLXBdistWheel,
+        },
+    )
+
+    package_data = {"mlx.core": ["*.pyi"]}
+
+    extras = {
+        "dev": [
+            "numpy>=2",
+            "pre-commit",
+            "psutil",
+            "torch>=2.9",
+            "typing_extensions",
+        ],
+    }
+    entry_points = {
+        "console_scripts": [
+            "mlx.launch = mlx._distributed_utils.launch:main",
+            "mlx.distributed_config = mlx._distributed_utils.config:main",
+        ]
+    }
+    install_requires = []
+
+    # Release builds for PyPi are in two stages.
+    # Each stage should be run from a clean build:
+    #   python setup.py clean --all
+    #
+    # Stage 1:
+    #  - Triggered with `MLX_BUILD_STAGE=1`
+    #  - Include everything except backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc)
+    #  - Wheel has Python ABI and platform tags
+    #  - Wheel should be built for the cross-product of python version and platforms
+    #  - Package name is mlx and it depends on subpackage in stage 2 (e.g. mlx-metal)
+    # Stage 2:
+    #  - Triggered with `MLX_BUILD_STAGE=2`
+    #  - Includes only backend-specific binaries (e.g. libmlx.so, mlx.metallib, etc)
+    #  - Wheel has only platform tags
+    #  - Wheel should be built only for different platforms
+    #  - Package name is back-end specific, e.g mlx-metal
+    if build_stage != 2:
+        if build_stage == 1:
+            install_requires.append(
+                f'mlx-metal=={version}; platform_system == "Darwin"'
+            )
+            extras["cuda"] = [f'mlx-cuda-12=={version}; platform_system == "Linux"']
+            for toolkit in [12, 13]:
+                extras[f"cuda{toolkit}"] = [
+                    f'mlx-cuda-{toolkit}=={version}; platform_system == "Linux"'
+                ]
+            extras["cpu"] = [f'mlx-cpu=={version}; platform_system == "Linux"']
+
+        _setup(
+            name="mlx",
+            packages=packages,
+            extras_require=extras,
+            entry_points=entry_points,
+            install_requires=install_requires,
+            package_data=package_data,
+        )
+    else:
+        if build_macos:
+            name = "mlx-metal"
+        elif build_cuda:
+            toolkit = cuda_toolkit_major_version()
+            name = f"mlx-cuda-{toolkit}"
+            # Note: update following files when new dependency is added:
+            # * .github/actions/build-cuda-release/action.yml
+            # * mlx/backend/cuda/CMakeLists.txt
+            if toolkit == 12:
+                install_requires += [
+                    "nvidia-cublas-cu12==12.9.*",
+                    "nvidia-cuda-nvrtc-cu12==12.9.*",
+                ]
+            elif toolkit == 13:
+                install_requires += [
+                    "nvidia-cublas",
+                    "nvidia-cuda-nvrtc",
+                ]
+            else:
+                raise ValueError(f"Unknown toolkit {toolkit}")
+            install_requires += [
+                f"nvidia-cudnn-cu{toolkit}==9.*",
+                f"nvidia-nccl-cu{toolkit}",
+            ]
+
+        else:
+            name = "mlx-cpu"
+        _setup(
+            name=name,
+            packages=["mlx"],
+            install_requires=install_requires,
+        )

data/mlx/tests/CMakeLists.txt

@@ -0,0 +1,44 @@
+FetchContent_Declare(
+  doctest
+  GIT_REPOSITORY https://github.com/onqtam/doctest.git
+  GIT_TAG v2.4.12)
+FetchContent_MakeAvailable(doctest)
+
+add_executable(tests ${PROJECT_SOURCE_DIR}/tests/tests.cpp)
+
+if(MLX_BUILD_METAL OR MLX_BUILD_CUDA)
+  set(METAL_TEST_SOURCES gpu_tests.cpp)
+endif()
+
+include(${doctest_SOURCE_DIR}/scripts/cmake/doctest.cmake)
+
+target_sources(
+  tests
+  PRIVATE allocator_tests.cpp
+          array_tests.cpp
+          arg_reduce_tests.cpp
+          autograd_tests.cpp
+          blas_tests.cpp
+          compile_tests.cpp
+          custom_vjp_tests.cpp
+          creations_tests.cpp
+          device_tests.cpp
+          einsum_tests.cpp
+          export_import_tests.cpp
+          eval_tests.cpp
+          fft_tests.cpp
+          load_tests.cpp
+          ops_tests.cpp
+          random_tests.cpp
+          scheduler_tests.cpp
+          utils_tests.cpp
+          vmap_tests.cpp
+          linalg_tests.cpp
+          ${METAL_TEST_SOURCES})
+
+target_link_libraries(tests PRIVATE mlx doctest)
+target_compile_options(tests PRIVATE ${SANITIZER_COMPILE_FLAGS})
+target_link_options(tests PRIVATE ${SANITIZER_LINK_FLAGS})
+
+doctest_discover_tests(tests)
+add_test(NAME tests COMMAND tests)

data/mlx/tests/allocator_tests.cpp

@@ -0,0 +1,41 @@
+// Copyright © 2023 Apple Inc.
+
+#include <stdexcept>
+
+#include "doctest/doctest.h"
+
+#include "mlx/allocator.h"
+
+using namespace mlx::core;
+
+TEST_CASE("test simple allocations") {
+  {
+    auto buffer = allocator::malloc(sizeof(float));
+    auto fptr = static_cast<float*>(buffer.raw_ptr());
+    *fptr = 0.5f;
+    CHECK_EQ(*fptr, 0.5f);
+    allocator::free(buffer);
+  }
+
+  {
+    auto buffer = allocator::malloc(128 * sizeof(int));
+    int* ptr = static_cast<int*>(buffer.raw_ptr());
+    for (int i = 0; i < 128; ++i) {
+      ptr[i] = i;
+    }
+    allocator::free(buffer);
+  }
+
+  {
+    auto buffer = allocator::malloc(0);
+    allocator::free(buffer);
+  }
+}
+
+TEST_CASE("test large allocations") {
+  size_t size = 1 << 30;
+  for (int i = 0; i < 100; ++i) {
+    auto buffer = allocator::malloc(size);
+    allocator::free(buffer);
+  }
+}

data/mlx/tests/arg_reduce_tests.cpp

@@ -0,0 +1,204 @@
+// Copyright © 2023 Apple Inc.
+
+#include "doctest/doctest.h"
+
+#include "mlx/mlx.h"
+#include "mlx/primitives.h"
+
+using namespace mlx::core;
+
+void test_arg_reduce_small(
+    Device d,
+    const array& x,
+    ArgReduce::ReduceType r,
+    Shape out_shape,
+    int axis,
+    std::vector<int> expected_output) {
+  auto s = default_stream(d);
+  auto y =
+      array(out_shape, uint32, std::make_shared<ArgReduce>(s, r, axis), {x});
+  y.eval();
+  const uint32_t* ydata = y.data<uint32_t>();
+  for (int i = 0; i < y.size(); i++) {
+    CHECK_EQ(expected_output[i], ydata[i]);
+  }
+}
+
+void test_arg_reduce_against_cpu(
+    const array& x,
+    ArgReduce::ReduceType r,
+    Shape out_shape,
+    int axis) {
+  auto y1 = array(
+      out_shape,
+      uint32,
+      std::make_shared<ArgReduce>(default_stream(Device::cpu), r, axis),
+      {x});
+  auto y2 = array(
+      out_shape,
+      uint32,
+      std::make_shared<ArgReduce>(default_stream(Device::gpu), r, axis),
+      {x});
+  y1.eval();
+  y2.eval();
+  CHECK(array_equal(y1, y2).item<bool>());
+}
+
+TEST_CASE("test arg reduce small") {
+  auto x = array(
+      {0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5,
+       0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5},
+      {2, 3, 4});
+  test_arg_reduce_small(
+      Device::cpu, x, ArgReduce::ArgMin, {2, 3}, 2, {0, 1, 3, 0, 1, 3});
+  test_arg_reduce_small(
+      Device::cpu, x, ArgReduce::ArgMin, {2, 4}, 1, {0, 1, 1, 2, 0, 1, 1, 2});
+  test_arg_reduce_small(
+      Device::cpu,
+      x,
+      ArgReduce::ArgMin,
+      {3, 4},
+      0,
+      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+  test_arg_reduce_small(
+      Device::cpu, x, ArgReduce::ArgMax, {2, 3}, 2, {3, 0, 1, 3, 0, 1});
+  test_arg_reduce_small(
+      Device::cpu, x, ArgReduce::ArgMax, {2, 4}, 1, {1, 2, 2, 0, 1, 2, 2, 0});
+  test_arg_reduce_small(
+      Device::cpu,
+      x,
+      ArgReduce::ArgMax,
+      {3, 4},
+      0,
+      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+
+  if (!metal::is_available()) {
+    INFO("Skipping arg reduction gpu tests");
+    return;
+  }
+
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMin, {2, 3}, 2, {0, 1, 3, 0, 1, 3});
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMin, {2, 4}, 1, {0, 1, 1, 2, 0, 1, 1, 2});
+  test_arg_reduce_small(
+      Device::gpu,
+      x,
+      ArgReduce::ArgMin,
+      {3, 4},
+      0,
+      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMax, {2, 3}, 2, {3, 0, 1, 3, 0, 1});
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMax, {2, 4}, 1, {1, 2, 2, 0, 1, 2, 2, 0});
+  test_arg_reduce_small(
+      Device::gpu,
+      x,
+      ArgReduce::ArgMax,
+      {3, 4},
+      0,
+      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+}
+
+TEST_CASE("test arg reduce against cpu") {
+  if (!metal::is_available()) {
+    INFO("Skipping arg reduction gpu tests");
+    return;
+  }
+
+  auto x = random::uniform(array(0.0), array(1.0), {127, 92, 55});
+  x.eval();
+  test_arg_reduce_against_cpu(x, ArgReduce::ArgMin, {127, 92}, 2);
+  test_arg_reduce_against_cpu(x, ArgReduce::ArgMin, {127, 55}, 1);
+  test_arg_reduce_against_cpu(x, ArgReduce::ArgMin, {92, 55}, 0);
+  test_arg_reduce_against_cpu(x, ArgReduce::ArgMax, {127, 92}, 2);
+  test_arg_reduce_against_cpu(x, ArgReduce::ArgMax, {127, 55}, 1);
+  test_arg_reduce_against_cpu(x, ArgReduce::ArgMax, {92, 55}, 0);
+
+  auto y = random::uniform(array(0.0), array(1.0), {1234});
+  y.eval();
+  test_arg_reduce_against_cpu(y, ArgReduce::ArgMin, {}, 0);
+  test_arg_reduce_against_cpu(y, ArgReduce::ArgMax, {}, 0);
+}
+
+void test_arg_reduce_small_bool(
+    Device d,
+    ArgReduce::ReduceType r,
+    Shape out_shape,
+    int axis,
+    std::vector<int> expected_output) {
+  auto s = default_stream(d);
+  auto x = array(
+      {0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5,
+       0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5},
+      {2, 3, 4});
+  x.eval();
+  auto y =
+      array(out_shape, uint32, std::make_shared<ArgReduce>(s, r, axis), {x});
+  y.eval();
+  const uint32_t* ydata = y.data<uint32_t>();
+  for (int i = 0; i < y.size(); i++) {
+    CHECK_EQ(expected_output[i], ydata[i]);
+  }
+}
+
+TEST_CASE("test arg reduce bool") {
+  if (!metal::is_available()) {
+    INFO("Skipping arg reduction gpu tests");
+    return;
+  }
+  auto x = array(
+      {false, true,  true,  false, false, false, false, true,
+       true,  false, true,  true,  false, true,  true,  false,
+       false, false, false, true,  true,  false, true,  true},
+      {2, 3, 4});
+  x.eval();
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMin, {2, 3}, 2, {0, 0, 1, 0, 0, 1});
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMin, {2, 4}, 1, {0, 1, 1, 0, 0, 1, 1, 0});
+  test_arg_reduce_small(
+      Device::gpu,
+      x,
+      ArgReduce::ArgMin,
+      {3, 4},
+      0,
+      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMax, {2, 3}, 2, {1, 3, 0, 1, 3, 0});
+  test_arg_reduce_small(
+      Device::gpu, x, ArgReduce::ArgMax, {2, 4}, 1, {2, 0, 0, 1, 2, 0, 0, 1});
+  test_arg_reduce_small(
+      Device::gpu,
+      x,
+      ArgReduce::ArgMax,
+      {3, 4},
+      0,
+      {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+}
+
+TEST_CASE("test arg reduce edge cases") {
+  auto a = argmin(array(1.0));
+  CHECK_EQ(a.item<uint32_t>(), 0);
+  auto b = argmax(array(1.0));
+  CHECK_EQ(b.item<uint32_t>(), 0);
+  CHECK_THROWS(argmin(array({})));
+  CHECK_THROWS(argmax(array({})));
+}
+
+TEST_CASE("test arg reduce irregular strides") {
+  auto x = array(
+      {0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5,
+       0, 2, 1, 7, 5, -5, 0, 2, 1, 7, 5, -5},
+      {2, 3, 4});
+  x = transpose(x, {2, 0, 1});
+  x.eval();
+  test_arg_reduce_small(
+      Device::cpu, x, ArgReduce::ArgMin, {4, 2}, 2, {0, 0, 1, 1, 1, 1, 2, 2});
+
+  if (!metal::is_available()) {
+    INFO("Skipping arg reduction gpu tests");
+    return;
+  }
+}