diff --git a/CMakeLists.txt b/CMakeLists.txt index fcf4ea77..464eb54c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -35,7 +35,7 @@ else() FetchContent_Declare( mlx GIT_REPOSITORY "https://github.com/ml-explore/mlx.git" - GIT_TAG v0.31.2) + GIT_TAG v0.32.0) FetchContent_MakeAvailable(mlx) endif() diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 7cbacf31..dd4edbf4 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -26,3 +26,6 @@ target_link_libraries(example-gguf PUBLIC mlxc) add_executable(example-graph ${CMAKE_CURRENT_LIST_DIR}/example-graph.c) target_link_libraries(example-graph PUBLIC mlxc) + +add_executable(example-ops ${CMAKE_CURRENT_LIST_DIR}/example-ops.c) +target_link_libraries(example-ops PUBLIC mlxc) diff --git a/examples/example-ops.c b/examples/example-ops.c new file mode 100644 index 00000000..0f6a89db --- /dev/null +++ b/examples/example-ops.c @@ -0,0 +1,77 @@ +/* Copyright © 2023-2024 Apple Inc. */ + +#include +#include "mlx/c/mlx.h" + +void print_array(const char* msg, mlx_array arr) { + mlx_string str = mlx_string_new(); + mlx_array_tostring(&str, arr); + printf("%s\n%s\n", msg, mlx_string_data(str)); + mlx_string_free(str); +} + +/* Exercises a range of array, elementwise, and linalg ops. */ +int main(void) { + mlx_stream stream = mlx_default_cpu_stream_new(); + + float data[] = {1, -2, 0, 4, -5, 6}; + int shape[] = {2, 3}; + mlx_array arr = mlx_array_new_data(data, shape, 2, MLX_FLOAT32); + mlx_array res = mlx_array_new(); + + mlx_positive(&res, arr, stream); + print_array("positive", res); + + mlx_trunc(&res, arr, stream); + print_array("trunc", res); + + mlx_count_nonzero(&res, arr, 1, false, stream); + print_array("count_nonzero (axis=1)", res); + + mlx_diff(&res, arr, 1, 1, stream); + print_array("diff (n=1, axis=1)", res); + + int axes[] = {1}; + mlx_flip(&res, arr, axes, 1, stream); + print_array("flip (axis=1)", res); + + float rhs_data[] = {6, 5, 4, 3, 2, 1}; + mlx_array rhs = mlx_array_new_data(rhs_data, shape, 2, MLX_FLOAT32); + mlx_vecdot(&res, arr, rhs, 1, stream); + print_array("vecdot (axis=1)", res); + + mlx_array lbool = mlx_array_new_bool(true); + mlx_array rbool = mlx_array_new_bool(false); + mlx_logical_xor(&res, lbool, rbool, stream); + print_array("logical_xor", res); + + mlx_vector_array parts = mlx_vector_array_new(); + mlx_unstack(&parts, arr, 0, stream); + printf("unstack (axis=0): %zu arrays\n", mlx_vector_array_size(parts)); + + float square_data[] = {4, 3, 6, 3}; + int square_shape[] = {2, 2}; + mlx_array square = + mlx_array_new_data(square_data, square_shape, 2, MLX_FLOAT32); + mlx_linalg_det(&res, square, stream); + print_array("linalg_det", res); + + mlx_array sign = mlx_array_new(); + mlx_array logabsdet = mlx_array_new(); + mlx_linalg_slogdet(&sign, &logabsdet, square, stream); + print_array("slogdet sign", sign); + print_array("slogdet logabsdet", logabsdet); + + mlx_array_free(arr); + mlx_array_free(rhs); + mlx_array_free(res); + mlx_array_free(lbool); + mlx_array_free(rbool); + mlx_array_free(square); + mlx_array_free(sign); + mlx_array_free(logabsdet); + mlx_vector_array_free(parts); + mlx_stream_free(stream); + + return 0; +} diff --git a/examples/example.c b/examples/example.c index 2e803960..0d90a532 100644 --- a/examples/example.c +++ b/examples/example.c @@ -58,6 +58,11 @@ void gpu_info(void) { mlx_device_info_free(info); mlx_device_free(dev); + const char* metallib_path = NULL; + if (mlx_metal_get_metallib_path(&metallib_path) == 0) + printf( + "Metal library path: %s\n", metallib_path ? metallib_path : "(null)"); + printf("==================================================\n"); } diff --git a/mlx/c/linalg.cpp b/mlx/c/linalg.cpp index a9a3c43c..4ea04a8b 100644 --- a/mlx/c/linalg.cpp +++ b/mlx/c/linalg.cpp @@ -57,6 +57,17 @@ extern "C" int mlx_linalg_cross( } return 0; } +extern "C" int +mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::linalg::det(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -261,6 +272,24 @@ extern "C" int mlx_linalg_qr( } return 0; } +extern "C" int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s) { + try { + { + auto [tpl_0, tpl_1] = + mlx::core::linalg::slogdet(mlx_array_get_(a), mlx_stream_get_(s)); + mlx_array_set_(*res_0, tpl_0); + mlx_array_set_(*res_1, tpl_1); + }; + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_linalg_solve( mlx_array* res, const mlx_array a, diff --git a/mlx/c/linalg.h b/mlx/c/linalg.h index 91d5d661..6d5b5ae1 100644 --- a/mlx/c/linalg.h +++ b/mlx/c/linalg.h @@ -44,6 +44,7 @@ int mlx_linalg_cross( const mlx_array b, int axis, const mlx_stream s); +int mlx_linalg_det(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_linalg_eig( mlx_array* res_0, mlx_array* res_1, @@ -97,6 +98,11 @@ int mlx_linalg_qr( mlx_array* res_1, const mlx_array a, const mlx_stream s); +int mlx_linalg_slogdet( + mlx_array* res_0, + mlx_array* res_1, + const mlx_array a, + const mlx_stream s); int mlx_linalg_solve( mlx_array* res, const mlx_array a, diff --git a/mlx/c/metal.cpp b/mlx/c/metal.cpp index 46d8034d..c93a0e51 100644 --- a/mlx/c/metal.cpp +++ b/mlx/c/metal.cpp @@ -8,6 +8,15 @@ #include "mlx/c/error.h" #include "mlx/c/private/mlx.h" +extern "C" int mlx_metal_get_metallib_path(const char** res) { + try { + *res = mlx::core::metal::get_metallib_path().c_str(); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_metal_is_available(bool* res) { try { *res = mlx::core::metal::is_available(); @@ -17,6 +26,15 @@ extern "C" int mlx_metal_is_available(bool* res) { } return 0; } +extern "C" int mlx_metal_set_metallib_path(const char* path) { + try { + mlx::core::metal::set_metallib_path(std::string(path)); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_metal_start_capture(const char* path) { try { mlx::core::metal::start_capture(std::string(path)); diff --git a/mlx/c/metal.h b/mlx/c/metal.h index 5877b224..a1bcf553 100644 --- a/mlx/c/metal.h +++ b/mlx/c/metal.h @@ -28,7 +28,9 @@ extern "C" { */ /**@{*/ +int mlx_metal_get_metallib_path(const char** res); int mlx_metal_is_available(bool* res); +int mlx_metal_set_metallib_path(const char* path); int mlx_metal_start_capture(const char* path); int mlx_metal_stop_capture(void); diff --git a/mlx/c/ops.cpp b/mlx/c/ops.cpp index 40b5a0cf..999ee166 100644 --- a/mlx/c/ops.cpp +++ b/mlx/c/ops.cpp @@ -446,12 +446,17 @@ extern "C" int mlx_astype( mlx_array* res, const mlx_array a, mlx_dtype dtype, + mlx_optional_bool copy, const mlx_stream s) { try { mlx_array_set_( *res, mlx::core::astype( - mlx_array_get_(a), mlx_dtype_to_cpp(dtype), mlx_stream_get_(s))); + mlx_array_get_(a), + mlx_dtype_to_cpp(dtype), + (copy.has_value ? std::make_optional(copy.value) + : std::nullopt), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -973,6 +978,23 @@ extern "C" int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s) { } return 0; } +extern "C" int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::count_nonzero( + mlx_array_get_(a), axis, keepdims, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_cummax( mlx_array* res, const mlx_array a, @@ -1015,12 +1037,20 @@ extern "C" int mlx_cumprod( int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) { try { mlx_array_set_( *res, mlx::core::cumprod( - mlx_array_get_(a), axis, reverse, inclusive, mlx_stream_get_(s))); + mlx_array_get_(a), + axis, + reverse, + inclusive, + (dtype.has_value ? std::make_optional( + mlx_dtype_to_cpp(dtype.value)) + : std::nullopt), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -1033,12 +1063,20 @@ extern "C" int mlx_cumsum( int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s) { try { mlx_array_set_( *res, mlx::core::cumsum( - mlx_array_get_(a), axis, reverse, inclusive, mlx_stream_get_(s))); + mlx_array_get_(a), + axis, + reverse, + inclusive, + (dtype.has_value ? std::make_optional( + mlx_dtype_to_cpp(dtype.value)) + : std::nullopt), + mlx_stream_get_(s))); } catch (std::exception& e) { mlx_error(e.what()); return 1; @@ -1137,6 +1175,21 @@ extern "C" int mlx_diagonal( } return 0; } +extern "C" int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::diff(mlx_array_get_(a), n, axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_divide( mlx_array* res, const mlx_array a, @@ -1311,6 +1364,25 @@ extern "C" int mlx_flatten( } return 0; } +extern "C" int mlx_flip( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::flip( + mlx_array_get_(a), + std::vector(axes, axes + axes_num), + mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s) { try { @@ -1895,6 +1967,22 @@ extern "C" int mlx_logical_or( } return 0; } +extern "C" int mlx_logical_xor( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::logical_xor( + mlx_array_get_(a), mlx_array_get_(b), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_logsumexp_axes( mlx_array* res, const mlx_array a, @@ -2432,6 +2520,17 @@ mlx_partition(mlx_array* res, const mlx_array a, int kth, const mlx_stream s) { } return 0; } +extern "C" int +mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::positive(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_power( mlx_array* res, const mlx_array a, @@ -3948,6 +4047,17 @@ mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s) { } return 0; } +extern "C" int +mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s) { + try { + mlx_array_set_( + *res, mlx::core::trunc(mlx_array_get_(a), mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -3969,6 +4079,20 @@ extern "C" int mlx_unflatten( } return 0; } +extern "C" int mlx_unstack( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s) { + try { + mlx_vector_array_set_( + *res, mlx::core::unstack(mlx_array_get_(a), axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -4026,6 +4150,23 @@ extern "C" int mlx_var( } return 0; } +extern "C" int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s) { + try { + mlx_array_set_( + *res, + mlx::core::vecdot( + mlx_array_get_(a), mlx_array_get_(b), axis, mlx_stream_get_(s))); + } catch (std::exception& e) { + mlx_error(e.what()); + return 1; + } + return 0; +} extern "C" int mlx_view( mlx_array* res, const mlx_array a, diff --git a/mlx/c/ops.h b/mlx/c/ops.h index 44fc09c4..9b013ea5 100644 --- a/mlx/c/ops.h +++ b/mlx/c/ops.h @@ -162,6 +162,7 @@ int mlx_astype( mlx_array* res, const mlx_array a, mlx_dtype dtype, + mlx_optional_bool copy, const mlx_stream s); int mlx_atleast_1d(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_atleast_2d(mlx_array* res, const mlx_array a, const mlx_stream s); @@ -323,6 +324,12 @@ int mlx_conv_transpose3d( int mlx_copy(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cos(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_cosh(mlx_array* res, const mlx_array a, const mlx_stream s); +int mlx_count_nonzero( + mlx_array* res, + const mlx_array a, + int axis, + bool keepdims, + const mlx_stream s); int mlx_cummax( mlx_array* res, const mlx_array a, @@ -343,6 +350,7 @@ int mlx_cumprod( int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_cumsum( mlx_array* res, @@ -350,6 +358,7 @@ int mlx_cumsum( int axis, bool reverse, bool inclusive, + mlx_optional_dtype dtype, const mlx_stream s); int mlx_degrees(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_depends( @@ -375,6 +384,12 @@ int mlx_diagonal( int axis1, int axis2, const mlx_stream s); +int mlx_diff( + mlx_array* res, + const mlx_array a, + int n, + int axis, + const mlx_stream s); int mlx_divide( mlx_array* res, const mlx_array a, @@ -423,6 +438,12 @@ int mlx_flatten( int start_axis, int end_axis, const mlx_stream s); +int mlx_flip( + mlx_array* res, + const mlx_array a, + const int* axes, + size_t axes_num, + const mlx_stream s); int mlx_floor(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_floor_divide( mlx_array* res, @@ -577,6 +598,11 @@ int mlx_logical_or( const mlx_array a, const mlx_array b, const mlx_stream s); +int mlx_logical_xor( + mlx_array* res, + const mlx_array a, + const mlx_array b, + const mlx_stream s); int mlx_logsumexp_axes( mlx_array* res, const mlx_array a, @@ -757,6 +783,7 @@ int mlx_partition( const mlx_array a, int kth, const mlx_stream s); +int mlx_positive(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_power( mlx_array* res, const mlx_array a, @@ -1231,6 +1258,7 @@ int mlx_tri( const mlx_stream s); int mlx_tril(mlx_array* res, const mlx_array x, int k, const mlx_stream s); int mlx_triu(mlx_array* res, const mlx_array x, int k, const mlx_stream s); +int mlx_trunc(mlx_array* res, const mlx_array a, const mlx_stream s); int mlx_unflatten( mlx_array* res, const mlx_array a, @@ -1238,6 +1266,11 @@ int mlx_unflatten( const int* shape, size_t shape_num, const mlx_stream s); +int mlx_unstack( + mlx_vector_array* res, + const mlx_array a, + int axis, + const mlx_stream s); int mlx_var_axes( mlx_array* res, const mlx_array a, @@ -1259,6 +1292,12 @@ int mlx_var( bool keepdims, int ddof, const mlx_stream s); +int mlx_vecdot( + mlx_array* res, + const mlx_array a, + const mlx_array b, + int axis, + const mlx_stream s); int mlx_view( mlx_array* res, const mlx_array a, diff --git a/mlx/c/optional.h b/mlx/c/optional.h index ff9ea14e..6bd93315 100644 --- a/mlx/c/optional.h +++ b/mlx/c/optional.h @@ -18,6 +18,14 @@ extern "C" { */ /**@{*/ +/** + * A bool optional. + */ +typedef struct mlx_optional_bool_ { + bool value; + bool has_value; +} mlx_optional_bool; + /** * A int optional. */ diff --git a/python/generator.py b/python/generator.py index 88b0450b..da18b77d 100644 --- a/python/generator.py +++ b/python/generator.py @@ -82,6 +82,8 @@ def preprocess_header(content): """Simple preprocessor that strips MLX_API macro without resolving includes.""" # Remove MLX_API macro (appears as MLX_API or MLX_API followed by space) content = re.sub(r"\bMLX_API\s*", "", content) + # Strip storage specifiers the parser can't handle in class member decls + content = re.sub(r"\bthread_local\s*", "", content) return content diff --git a/python/mlxtypes.py b/python/mlxtypes.py index 66f52d7c..1220dc2a 100644 --- a/python/mlxtypes.py +++ b/python/mlxtypes.py @@ -438,10 +438,10 @@ def register_return_tuple_type(cpp_types, alts=[]): "alt": "std::string", "c_to_cpp": lambda s: "std::string(" + s + ")", "c_arg": lambda s, untyped=False: s if untyped else "const char* " + s, - "c_return_arg": lambda s, untyped=False: s if untyped else "char** " + s, + "c_return_arg": lambda s, untyped=False: s if untyped else "const char** " + s, # "c_new": lambda s: "char* " + s, # "free": lambda s: "", - "c_assign_from_cpp": lambda d, s: d + " = " + s + ".c_str()", + "c_assign_from_cpp": lambda d, s: "*" + d + " = " + s + ".c_str()", } ) @@ -489,7 +489,7 @@ def register_return_tuple_type(cpp_types, alts=[]): ) types[-1]["alt"] = "std::uintptr_t" -for cpptype in ["float", "int", "mlx::core::Dtype"]: +for cpptype in ["bool", "float", "int", "mlx::core::Dtype"]: typedef = find_cpp_type(cpptype) ctype = typedef["c"] alttype = typedef["alt"]