update to mlx v0.25.1 (#217)

- adopt mlx v0.25.1 and mlx-c v0.2.0

- add saveToData and load from data (safetensors) -- fix #214
- fixes #211 (moves metalKernel API back to what it was)

- add withError and withErrorHandler error handling
- add function import/export: https://ml-explore.github.io/mlx/build/html/usage/export.html
 
- use an evalLock for eval, asyncEval and Stream creation

Author: davidkoski

CMakeLists.txt

@@ -11,7 +11,7 @@ endif()
 FetchContent_Declare(
   mlx-c
   GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git"
-  GIT_TAG "v0.1.2")
+  GIT_TAG "v0.2.0")
 FetchContent_MakeAvailable(mlx-c)
 
 # swift-numerics

Package.swift

@@ -87,12 +87,11 @@ let package = Package(
 
                 // do not build distributed support (yet)
                 "mlx/mlx/distributed/mpi/mpi.cpp",
-                "mlx/mlx/distributed/ops.cpp",
-                "mlx/mlx/distributed/primitives.cpp",
                 "mlx/mlx/distributed/ring/ring.cpp",
 
-                "mlx/mlx/backend/cpu/gemms/no_bf16.cpp",
-                "mlx/mlx/backend/cpu/gemms/no_fp16.cpp",
+                // bnns instead of simd (accelerate)
+                "mlx/mlx/backend/cpu/gemms/simd_fp16.cpp",
+                "mlx/mlx/backend/cpu/gemms/simd_bf16.cpp",
             ],
 
             cSettings: [
@@ -112,6 +111,7 @@ let package = Package(
                 .define("_METAL_"),
                 .define("SWIFTPM_BUNDLE", to: "\"mlx-swift_Cmlx\""),
                 .define("METAL_PATH", to: "\"default.metallib\""),
+                .define("MLX_VERSION", to: "\"0.24.2\""),
             ],
             linkerSettings: [
                 .linkedFramework("Foundation"),

Source/Cmlx/include/mlx/c/closure.h

@@ -7,6 +7,7 @@
 #define MLX_CLOSURE_H
 
 #include "mlx/c/array.h"
+#include "mlx/c/map.h"
 #include "mlx/c/optional.h"
 #include "mlx/c/stream.h"
 #include "mlx/c/vector.h"
@@ -40,6 +41,32 @@ int mlx_closure_apply(
 
 mlx_closure mlx_closure_new_unary(int (*fun)(mlx_array*, const mlx_array));
 
+typedef struct mlx_closure_kwargs_ {
+  void* ctx;
+} mlx_closure_kwargs;
+mlx_closure_kwargs mlx_closure_kwargs_new();
+int mlx_closure_kwargs_free(mlx_closure_kwargs cls);
+mlx_closure_kwargs mlx_closure_kwargs_new_func(int (*fun)(
+    mlx_vector_array*,
+    const mlx_vector_array,
+    const mlx_map_string_to_array));
+mlx_closure_kwargs mlx_closure_kwargs_new_func_payload(
+    int (*fun)(
+        mlx_vector_array*,
+        const mlx_vector_array,
+        const mlx_map_string_to_array,
+        void*),
+    void* payload,
+    void (*dtor)(void*));
+int mlx_closure_kwargs_set(
+    mlx_closure_kwargs* cls,
+    const mlx_closure_kwargs src);
+int mlx_closure_kwargs_apply(
+    mlx_vector_array* res,
+    mlx_closure_kwargs cls,
+    const mlx_vector_array input_0,
+    const mlx_map_string_to_array input_1);
+
 typedef struct mlx_closure_value_and_grad_ {
   void* ctx;
 } mlx_closure_value_and_grad;

Source/Cmlx/include/mlx/c/compile.h

@@ -6,12 +6,14 @@
 #ifndef MLX_COMPILE_H
 #define MLX_COMPILE_H
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 
 #include "mlx/c/array.h"
 #include "mlx/c/closure.h"
 #include "mlx/c/distributed_group.h"
+#include "mlx/c/io_types.h"
 #include "mlx/c/map.h"
 #include "mlx/c/stream.h"
 #include "mlx/c/string.h"

Source/Cmlx/include/mlx/c/device.h

@@ -3,6 +3,8 @@
 #ifndef MLX_DEVICE_H
 #define MLX_DEVICE_H
 
+#include <stdbool.h>
+
 #include "mlx/c/string.h"
 
 #ifdef __cplusplus

Source/Cmlx/include/mlx/c/distributed.h

@@ -6,12 +6,14 @@
 #ifndef MLX_DISTRIBUTED_H
 #define MLX_DISTRIBUTED_H
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 
 #include "mlx/c/array.h"
 #include "mlx/c/closure.h"
 #include "mlx/c/distributed_group.h"
+#include "mlx/c/io_types.h"
 #include "mlx/c/map.h"
 #include "mlx/c/stream.h"
 #include "mlx/c/string.h"
@@ -30,6 +32,16 @@ int mlx_distributed_all_gather(
     const mlx_array x,
     const mlx_distributed_group group /* may be null */,
     const mlx_stream S);
+int mlx_distributed_all_max(
+    mlx_array* res,
+    const mlx_array x,
+    const mlx_distributed_group group /* may be null */,
+    const mlx_stream s);
+int mlx_distributed_all_min(
+    mlx_array* res,
+    const mlx_array x,
+    const mlx_distributed_group group /* may be null */,
+    const mlx_stream s);
 int mlx_distributed_all_sum(
     mlx_array* res,
     const mlx_array x,

Source/Cmlx/include/mlx/c/distributed_group.h

@@ -3,6 +3,8 @@
 #ifndef MLX_DISTRIBUTED_GROUP_H
 #define MLX_DISTRIBUTED_GROUP_H
 
+#include <stdbool.h>
+
 #include "mlx/c/stream.h"
 
 #ifdef __cplusplus

Source/Cmlx/include/mlx/c/export.h

@@ -0,0 +1,75 @@
+/* Copyright © 2023-2025 Apple Inc.                   */
+
+#ifndef MLX_EXPORT_H
+#define MLX_EXPORT_H
+
+#include <stdbool.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#include "mlx/c/array.h"
+#include "mlx/c/closure.h"
+#include "mlx/c/distributed_group.h"
+#include "mlx/c/io_types.h"
+#include "mlx/c/map.h"
+#include "mlx/c/stream.h"
+#include "mlx/c/string.h"
+#include "mlx/c/vector.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * \defgroup export Function serialization
+ */
+/**@{*/
+int mlx_export_function(
+    const char* file,
+    const mlx_closure fun,
+    const mlx_vector_array args,
+    bool shapeless);
+int mlx_export_function_kwargs(
+    const char* file,
+    const mlx_closure_kwargs fun,
+    const mlx_vector_array args,
+    const mlx_map_string_to_array kwargs,
+    bool shapeless);
+
+typedef struct mlx_function_exporter_ {
+  void* ctx;
+} mlx_function_exporter;
+mlx_function_exporter mlx_function_exporter_new(
+    const char* file,
+    const mlx_closure fun,
+    bool shapeless);
+int mlx_function_exporter_free(mlx_function_exporter xfunc);
+int mlx_function_exporter_apply(
+    const mlx_function_exporter xfunc,
+    const mlx_vector_array args);
+int mlx_function_exporter_apply_kwargs(
+    const mlx_function_exporter xfunc,
+    const mlx_vector_array args,
+    const mlx_map_string_to_array kwargs);
+
+typedef struct mlx_imported_function_ {
+  void* ctx;
+} mlx_imported_function;
+mlx_imported_function mlx_imported_function_new(const char* file);
+int mlx_imported_function_free(mlx_imported_function xfunc);
+int mlx_imported_function_apply(
+    mlx_vector_array* res,
+    const mlx_imported_function xfunc,
+    const mlx_vector_array args);
+int mlx_imported_function_apply_kwargs(
+    mlx_vector_array* res,
+    const mlx_imported_function xfunc,
+    const mlx_vector_array args,
+    const mlx_map_string_to_array kwargs);
+/**@}*/
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif

Source/Cmlx/include/mlx/c/fast.h

@@ -6,12 +6,14 @@
 #ifndef MLX_FAST_H
 #define MLX_FAST_H
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 
 #include "mlx/c/array.h"
 #include "mlx/c/closure.h"
 #include "mlx/c/distributed_group.h"
+#include "mlx/c/io_types.h"
 #include "mlx/c/map.h"
 #include "mlx/c/stream.h"
 #include "mlx/c/string.h"
@@ -49,77 +51,70 @@ int mlx_fast_layer_norm(
     float eps,
     const mlx_stream s);
 
-typedef struct mlx_fast_metal_kernel_ {
+typedef struct mlx_fast_metal_kernel_config_ {
   void* ctx;
-} mlx_fast_metal_kernel;
-mlx_fast_metal_kernel mlx_fast_metal_kernel_new(
-    const char* name,
-    const char* source,
-    const char* header);
-void mlx_fast_metal_kernel_free(mlx_fast_metal_kernel cls);
-int mlx_fast_metal_kernel_add_input_name(
-    mlx_fast_metal_kernel cls,
-    const char* name);
-int mlx_fast_metal_kernel_set_input_names(
-    mlx_fast_metal_kernel cls,
-    int num,
-    ...);
-int mlx_fast_metal_kernel_add_output_name(
-    mlx_fast_metal_kernel cls,
-    const char* name);
-int mlx_fast_metal_kernel_set_output_names(
-    mlx_fast_metal_kernel cls,
-    int num,
-    ...);
-int mlx_fast_metal_kernel_set_contiguous_rows(
-    mlx_fast_metal_kernel cls,
-    bool flag);
-int mlx_fast_metal_kernel_set_atomic_outputs(
-    mlx_fast_metal_kernel cls,
-    bool flag);
+} mlx_fast_metal_kernel_config;
+mlx_fast_metal_kernel_config mlx_fast_metal_kernel_config_new();
+void mlx_fast_metal_kernel_config_free(mlx_fast_metal_kernel_config cls);
 
-int mlx_fast_metal_kernel_add_output_arg(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_add_output_arg(
+    mlx_fast_metal_kernel_config cls,
     const int* shape,
     size_t size,
     mlx_dtype dtype);
-int mlx_fast_metal_kernel_set_grid(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_set_grid(
+    mlx_fast_metal_kernel_config cls,
     int grid1,
     int grid2,
     int grid3);
-int mlx_fast_metal_kernel_set_thread_group(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_set_thread_group(
+    mlx_fast_metal_kernel_config cls,
     int thread1,
     int thread2,
     int thread3);
-int mlx_fast_metal_kernel_set_init_value(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_set_init_value(
+    mlx_fast_metal_kernel_config cls,
     float value);
-int mlx_fast_metal_kernel_set_verbose(mlx_fast_metal_kernel cls, bool verbose);
-int mlx_fast_metal_kernel_add_template_arg_dtype(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_set_verbose(
+    mlx_fast_metal_kernel_config cls,
+    bool verbose);
+int mlx_fast_metal_kernel_config_add_template_arg_dtype(
+    mlx_fast_metal_kernel_config cls,
     const char* name,
     mlx_dtype dtype);
-int mlx_fast_metal_kernel_add_template_arg_int(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_add_template_arg_int(
+    mlx_fast_metal_kernel_config cls,
     const char* name,
     int value);
-int mlx_fast_metal_kernel_add_template_arg_bool(
-    mlx_fast_metal_kernel cls,
+int mlx_fast_metal_kernel_config_add_template_arg_bool(
+    mlx_fast_metal_kernel_config cls,
     const char* name,
     bool value);
 
+typedef struct mlx_fast_metal_kernel_ {
+  void* ctx;
+} mlx_fast_metal_kernel;
+
+mlx_fast_metal_kernel mlx_fast_metal_kernel_new(
+    const char* name,
+    const mlx_vector_string input_names,
+    const mlx_vector_string output_names,
+    const char* source,
+    const char* header,
+    bool ensure_row_contiguous,
+    bool atomic_outputs);
+void mlx_fast_metal_kernel_free(mlx_fast_metal_kernel cls);
 int mlx_fast_metal_kernel_apply(
     mlx_vector_array* outputs,
     mlx_fast_metal_kernel cls,
     const mlx_vector_array inputs,
+    const mlx_fast_metal_kernel_config config,
     const mlx_stream stream);
 
 int mlx_fast_rms_norm(
     mlx_array* res,
     const mlx_array x,
-    const mlx_array weight,
+    const mlx_array weight /* may be null */,
     float eps,
     const mlx_stream s);
 int mlx_fast_rope(
@@ -138,8 +133,8 @@ int mlx_fast_scaled_dot_product_attention(
     const mlx_array keys,
     const mlx_array values,
     float scale,
-    const mlx_array mask /* may be null */,
-    mlx_optional_int memory_efficient_threshold,
+    const char* mask_mode,
+    const mlx_vector_array mask_arrs,
     const mlx_stream s);
 /**@}*/
 

Source/Cmlx/include/mlx/c/fft.h

@@ -6,12 +6,14 @@
 #ifndef MLX_FFT_H
 #define MLX_FFT_H
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 
 #include "mlx/c/array.h"
 #include "mlx/c/closure.h"
 #include "mlx/c/distributed_group.h"
+#include "mlx/c/io_types.h"
 #include "mlx/c/map.h"
 #include "mlx/c/stream.h"
 #include "mlx/c/string.h"