mirror of
https://github.com/ml-explore/mlx.git
synced 2025-11-04 10:38:10 +08:00
c4230747a1
* redesign for faster cpu/gpu synch * load + more async CPU * use command encoder API and move more ops to use it * make fence back-end generic + CPU only fence * faster build * fix async eval * fixes + handle temporaries * fix / improve cpu conv * remove unused status, fix siblings * fix extensions * fix * fix no cpu build * format * comments * fix perf regression, remove unecessary abort * fix events, task limit cpu * fix waiting * fix donation / temporaries in normalization
79 lines
1.9 KiB
CMake
79 lines
1.9 KiB
CMake
cmake_minimum_required(VERSION 3.27)
|
|
|
|
project(_ext LANGUAGES CXX)
|
|
|
|
# ----------------------------- Setup -----------------------------
|
|
set(CMAKE_CXX_STANDARD 17)
|
|
set(CMAKE_CXX_STANDARD_REQUIRED ON)
|
|
set(CMAKE_POSITION_INDEPENDENT_CODE ON)
|
|
|
|
option(BUILD_SHARED_LIBS "Build extensions as a shared library" ON)
|
|
|
|
# ----------------------------- Dependencies -----------------------------
|
|
find_package(
|
|
Python 3.8
|
|
COMPONENTS Interpreter Development.Module
|
|
REQUIRED)
|
|
execute_process(
|
|
COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
|
|
OUTPUT_STRIP_TRAILING_WHITESPACE
|
|
OUTPUT_VARIABLE nanobind_ROOT)
|
|
find_package(nanobind CONFIG REQUIRED)
|
|
|
|
execute_process(
|
|
COMMAND "${Python_EXECUTABLE}" -m mlx --cmake-dir
|
|
OUTPUT_STRIP_TRAILING_WHITESPACE
|
|
OUTPUT_VARIABLE MLX_ROOT)
|
|
find_package(MLX CONFIG REQUIRED)
|
|
|
|
# ----------------------------- Extensions -----------------------------
|
|
|
|
# Add library
|
|
add_library(mlx_ext)
|
|
|
|
# Add sources
|
|
target_sources(mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.cpp)
|
|
|
|
# Add include headers
|
|
target_include_directories(mlx_ext PUBLIC ${CMAKE_CURRENT_LIST_DIR})
|
|
|
|
# Link to mlx
|
|
target_link_libraries(mlx_ext PUBLIC mlx)
|
|
|
|
# ----------------------------- Metal -----------------------------
|
|
|
|
# Build metallib
|
|
if(MLX_BUILD_METAL)
|
|
mlx_build_metallib(
|
|
TARGET
|
|
mlx_ext_metallib
|
|
TITLE
|
|
mlx_ext
|
|
SOURCES
|
|
${CMAKE_CURRENT_LIST_DIR}/axpby/axpby.metal
|
|
INCLUDE_DIRS
|
|
${PROJECT_SOURCE_DIR}
|
|
${MLX_INCLUDE_DIRS}
|
|
OUTPUT_DIRECTORY
|
|
${CMAKE_LIBRARY_OUTPUT_DIRECTORY})
|
|
|
|
add_dependencies(mlx_ext mlx_ext_metallib)
|
|
|
|
endif()
|
|
|
|
# ----------------------------- Python Bindings -----------------------------
|
|
nanobind_add_module(
|
|
_ext
|
|
NB_STATIC
|
|
STABLE_ABI
|
|
LTO
|
|
NOMINSIZE
|
|
NB_DOMAIN
|
|
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()
|