mirror of
https://github.com/ml-explore/mlx.git
synced 2025-07-23 10:02:12 +08:00
42afe27e12
* std and expm1 * actually add expm1 * fix linux * fix vjp * relax tol for linux test * Add it to the compilable primitives --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
115 lines
2.7 KiB
CMake
115 lines
2.7 KiB
CMake
set(
|
|
HEADERS
|
|
${CMAKE_CURRENT_SOURCE_DIR}/atomic.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/bf16.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/bf16_math.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/binary.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/complex.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/defines.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/erf.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/expm1f.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/indexing.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/unary.h
|
|
${CMAKE_CURRENT_SOURCE_DIR}/utils.h
|
|
)
|
|
|
|
set(
|
|
KERNELS
|
|
"arange"
|
|
"arg_reduce"
|
|
"binary"
|
|
"binary_two"
|
|
"conv"
|
|
"copy"
|
|
"gemv"
|
|
"quantized"
|
|
"random"
|
|
"rms_norm"
|
|
"layer_norm"
|
|
"rope"
|
|
"scan"
|
|
"scaled_dot_product_attention"
|
|
"softmax"
|
|
"sort"
|
|
"ternary"
|
|
"unary"
|
|
"gather"
|
|
"scatter"
|
|
)
|
|
|
|
function(build_kernel_base TARGET SRCFILE DEPS)
|
|
set(METAL_FLAGS -Wall -Wextra -fno-fast-math)
|
|
if(MLX_METAL_DEBUG)
|
|
set(METAL_FLAGS ${METAL_FLAGS}
|
|
-gline-tables-only
|
|
-frecord-sources)
|
|
endif()
|
|
add_custom_command(
|
|
COMMAND xcrun -sdk macosx metal
|
|
${METAL_FLAGS}
|
|
-c ${SRCFILE}
|
|
-I${PROJECT_SOURCE_DIR}
|
|
-o ${TARGET}.air
|
|
DEPENDS ${SRCFILE} ${DEPS}
|
|
OUTPUT ${TARGET}.air
|
|
COMMENT "Building ${TARGET}.air"
|
|
VERBATIM
|
|
)
|
|
endfunction(build_kernel_base)
|
|
|
|
function(build_kernel KERNEL)
|
|
set(SRCFILE ${CMAKE_CURRENT_SOURCE_DIR}/${KERNEL}.metal)
|
|
build_kernel_base(${KERNEL} ${SRCFILE} "${HEADERS}")
|
|
endfunction(build_kernel)
|
|
|
|
foreach(KERNEL ${KERNELS})
|
|
build_kernel(${KERNEL})
|
|
set(KERNEL_AIR ${KERNEL}.air ${KERNEL_AIR})
|
|
endforeach()
|
|
|
|
file(GLOB_RECURSE STEEL_KERNELS ${CMAKE_CURRENT_SOURCE_DIR}/steel/*.metal)
|
|
file(GLOB_RECURSE STEEL_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/steel/*.h)
|
|
|
|
foreach(KERNEL ${STEEL_KERNELS})
|
|
cmake_path(GET KERNEL STEM TARGET)
|
|
build_kernel_base(${TARGET} ${KERNEL} "${STEEL_HEADERS}")
|
|
set(KERNEL_AIR ${TARGET}.air ${KERNEL_AIR})
|
|
endforeach()
|
|
|
|
file(GLOB_RECURSE REDUCE_KERNELS ${CMAKE_CURRENT_SOURCE_DIR}/reduction/*.metal)
|
|
file(GLOB_RECURSE REDUCE_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/reduction/*.h)
|
|
|
|
foreach(KERNEL ${REDUCE_KERNELS})
|
|
cmake_path(GET KERNEL STEM TARGET)
|
|
build_kernel_base(${TARGET} ${KERNEL} "${REDUCE_HEADERS}")
|
|
set(KERNEL_AIR ${TARGET}.air ${KERNEL_AIR})
|
|
endforeach()
|
|
|
|
add_custom_command(
|
|
OUTPUT ${MLX_METAL_PATH}/mlx.metallib
|
|
COMMAND xcrun -sdk macosx metallib ${KERNEL_AIR} -o ${MLX_METAL_PATH}/mlx.metallib
|
|
DEPENDS ${KERNEL_AIR}
|
|
COMMENT "Building mlx.metallib"
|
|
VERBATIM
|
|
)
|
|
|
|
add_custom_target(
|
|
mlx-metallib
|
|
DEPENDS
|
|
${MLX_METAL_PATH}/mlx.metallib
|
|
)
|
|
|
|
add_dependencies(
|
|
mlx
|
|
mlx-metallib
|
|
)
|
|
|
|
# Install metallib
|
|
include(GNUInstallDirs)
|
|
|
|
install(
|
|
FILES ${MLX_METAL_PATH}/mlx.metallib
|
|
DESTINATION ${CMAKE_INSTALL_LIBDIR}
|
|
COMPONENT metallib
|
|
)
|