From f30b659291b1d279034c77891a2c700e50d7a6f3 Mon Sep 17 00:00:00 2001 From: Cheng Date: Wed, 27 Mar 2024 22:14:29 +0900 Subject: [PATCH] Make MLX build on x64 macOS (#901) The arm64 macbook pros are heavy and I usually care my intel one for mobile, it would be nice if I can play with MLX on it. To build with x64, user must pass `MLX_ENABLE_X64_MAC` to cmake: CMAKE_ARGS='-DMLX_ENABLE_X64_MAC=ON' python setup.py --- CMakeLists.txt | 60 ++++++++++++++++++++-------------- python/tests/mlx_tests.py | 4 +-- python/tests/test_fast_sdpa.py | 2 +- python/tests/test_ops.py | 4 +-- 4 files changed, 40 insertions(+), 30 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4b1ce18f3..5b54b1d0a 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -15,6 +15,7 @@ option(MLX_BUILD_EXAMPLES "Build examples for mlx" ON) option(MLX_BUILD_BENCHMARKS "Build benchmarks for mlx" OFF) option(MLX_BUILD_PYTHON_BINDINGS "Build python bindings for mlx" OFF) option(MLX_BUILD_METAL "Build metal backend" ON) +option(MLX_ENABLE_X64_MAC "Enable building for x64 macOS" OFF) option(BUILD_SHARED_LIBS "Build mlx as a shared library" OFF) if(NOT MLX_VERSION) @@ -23,23 +24,23 @@ endif() # --------------------- Processor tests ------------------------- -message(STATUS "Building MLX for ${CMAKE_HOST_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}") +message(STATUS "Building MLX for ${CMAKE_SYSTEM_PROCESSOR} processor on ${CMAKE_SYSTEM_NAME}") set(MLX_BUILD_ARM OFF) if (${CMAKE_SYSTEM_NAME} MATCHES "Darwin") - if (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64" AND ${CMAKE_HOST_APPLE}) - message(FATAL_ERROR - "Building for x86_64 on macOS is not supported." - " If you are on an Apple silicon system, check the build" - " documentation for possible fixes: " - "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source") - elseif (${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "x86_64") - message(WARNING - "Building for x86_64 on macOS is not supported." - " If you are on an Apple silicon system, " - " make sure you are building for arm64.") - elseif(${CMAKE_HOST_SYSTEM_PROCESSOR} MATCHES "arm64") + if(${CMAKE_SYSTEM_PROCESSOR} MATCHES "x86_64") + if(NOT MLX_ENABLE_X64_MAC) + message(FATAL_ERROR + "Building for x86_64 on macOS is not supported." + " If you are on an Apple silicon system, check the build" + " documentation for possible fixes: " + "https://ml-explore.github.io/mlx/build/html/install.html#build-from-source") + else() + message(WARNING "Building for x86_64 arch is not officially supported.") + endif() + set(MLX_BUILD_METAL OFF) + elseif(${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64") set(MLX_BUILD_ARM ON) endif() @@ -108,7 +109,27 @@ if (MLX_BUILD_ARM AND ACCELERATE_LIBRARY) else() message(STATUS "Accelerate or arm neon not found, using default backend.") set(MLX_BUILD_ACCELERATE OFF) - #set(BLA_VENDOR Generic) + if(${CMAKE_HOST_APPLE}) + # The blas shipped in macOS SDK is not supported, search homebrew for + # openblas instead. + set(BLA_VENDOR OpenBLAS) + set(LAPACK_ROOT "${LAPACK_ROOT};$ENV{LAPACK_ROOT};/usr/local/opt/openblas") + endif() + # Search and link with lapack. + find_package(LAPACK REQUIRED) + if (NOT LAPACK_FOUND) + message(FATAL_ERROR "Must have LAPACK installed") + endif() + find_path(LAPACK_INCLUDE_DIRS lapacke.h + /usr/include + /usr/local/include + /usr/local/opt/openblas/include) + message(STATUS "Lapack lib " ${LAPACK_LIBRARIES}) + message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS}) + target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS}) + target_link_libraries(mlx ${LAPACK_LIBRARIES}) + # List blas after lapack otherwise we may accidentally incldue an old version + # of lapack.h from the include dirs of blas. find_package(BLAS REQUIRED) if (NOT BLAS_FOUND) message(FATAL_ERROR "Must have BLAS installed") @@ -122,17 +143,6 @@ else() message(STATUS "Blas include " ${BLAS_INCLUDE_DIRS}) target_include_directories(mlx PRIVATE ${BLAS_INCLUDE_DIRS}) target_link_libraries(mlx ${BLAS_LIBRARIES}) - find_package(LAPACK REQUIRED) - if (NOT LAPACK_FOUND) - message(FATAL_ERROR "Must have LAPACK installed") - endif() - find_path(LAPACK_INCLUDE_DIRS lapacke.h - /usr/include - /usr/local/include) - message(STATUS "Lapack lib " ${LAPACK_LIBRARIES}) - message(STATUS "Lapack include " ${LAPACK_INCLUDE_DIRS}) - target_include_directories(mlx PRIVATE ${LAPACK_INCLUDE_DIRS}) - target_link_libraries(mlx ${LAPACK_LIBRARIES}) endif() add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/mlx) diff --git a/python/tests/mlx_tests.py b/python/tests/mlx_tests.py index 477e81c8c..f446b5e67 100644 --- a/python/tests/mlx_tests.py +++ b/python/tests/mlx_tests.py @@ -11,8 +11,8 @@ import numpy as np class MLXTestCase(unittest.TestCase): @property - def is_linux(self): - return platform.system() == "Linux" + def is_apple_silicon(self): + return platform.machine() == "arm64" and platform.system() == "Darwin" def setUp(self): self.default = mx.default_device() diff --git a/python/tests/test_fast_sdpa.py b/python/tests/test_fast_sdpa.py index 4be45a552..b4fa07395 100644 --- a/python/tests/test_fast_sdpa.py +++ b/python/tests/test_fast_sdpa.py @@ -62,7 +62,7 @@ class TestFastSDPA(mlx_tests.MLXTestCase): B = 1 H = 32 dtypes = [np.float32] - if not self.is_linux: + if self.is_apple_silicon: dtypes.append(np.half) for SEQUENCE_LENGTH in [1, 7, 9, 32, 63, 67, 129, 400, 2000]: diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index e17c5fce9..417272ad8 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1772,8 +1772,8 @@ class TestOps(mlx_tests.MLXTestCase): ) def test_tensordot(self): - # No fp16 matmuls on linux - if self.is_linux: + # No fp16 matmuls on common cpu backend + if not self.is_apple_silicon: dtypes = [mx.float32] else: dtypes = [mx.float16, mx.float32]