Update mlx-c to 0.1.2 (#204)

* update to mlx-c 0.1.2

* add linalg + kron + flatten

* format

* fix cmake build

* Add David's updates + better float64 support.

Co-authored-by: DavidKoski <46639364+davidkoski@users.noreply.github.com>

* remove unwraps

---------

Co-authored-by: Alex Barron <abarron22@apple.com>
Co-authored-by: DavidKoski <46639364+davidkoski@users.noreply.github.com>

Author: 3 people

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.0")
+  GIT_TAG "v0.1.2")
 FetchContent_MakeAvailable(mlx-c)
 
 # swift-numerics

Package.swift

@@ -51,9 +51,8 @@ let package = Package(
                 "fmt/src/fmt.cc",
 
                 // these are selected conditionally
-                // via mlx-conditional/compiled_conditional.cpp
-                "mlx/mlx/backend/common/compiled_nocpu.cpp",
-                "mlx/mlx/backend/common/compiled_cpu.cpp",
+                "mlx/mlx/backend/no_cpu/compiled.cpp",
+                "mlx/mlx/backend/cpu/compiled.cpp",
 
                 // mlx files that are not part of the build
                 "mlx/ACKNOWLEDGMENTS.md",
@@ -79,10 +78,6 @@ let package = Package(
 
                 "mlx/mlx/backend/common/default_primitives.cpp",
 
-                // this uses neon code and will not build on x86 (e.g. via Release).
-                // see mlx-conditional/accelerate-softmax.cpp
-                "mlx/mlx/backend/accelerate/softmax.cpp",
-
                 // build variants (we are opting _out_ of these)
                 "mlx/mlx/io/no_safetensors.cpp",
                 "mlx/mlx/io/gguf.cpp",
@@ -93,9 +88,13 @@ let package = Package(
                 "mlx/mlx/backend/metal/nojit_kernels.cpp",
 
                 // do not build distributed support (yet)
-                "mlx/mlx/distributed/mpi",
+                "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",
             ],
 
             cSettings: [
@@ -110,6 +109,7 @@ let package = Package(
                 .headerSearchPath("json/single_include/nlohmann"),
                 .headerSearchPath("fmt/include"),
 
+                .define("MLX_USE_ACCELERATE"),
                 .define("ACCELERATE_NEW_LAPACK"),
                 .define("_METAL_"),
                 .define("SWIFTPM_BUNDLE", to: "\"mlx-swift_Cmlx\""),

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

@@ -46,6 +46,7 @@ typedef enum mlx_dtype_ {
   MLX_INT64,
   MLX_FLOAT16,
   MLX_FLOAT32,
+  MLX_FLOAT64,
   MLX_BFLOAT16,
   MLX_COMPLEX64,
 } mlx_dtype;
@@ -78,10 +79,24 @@ mlx_array mlx_array_new_bool(bool val);
  * New array from a int scalar.
  */
 mlx_array mlx_array_new_int(int val);
+/**
+ * New array from a float32 scalar.
+ */
+mlx_array mlx_array_new_float32(float val);
 /**
  * New array from a float scalar.
+ * Same as float32.
  */
 mlx_array mlx_array_new_float(float val);
+/**
+ * New array from a float64 scalar.
+ */
+mlx_array mlx_array_new_float64(double val);
+/**
+ * New array from a double scalar.
+ * Same as float64.
+ */
+mlx_array mlx_array_new_double(double val);
 /**
  * New array from a complex scalar.
  */
@@ -110,10 +125,22 @@ int mlx_array_set_bool(mlx_array* arr, bool val);
  * Set array to a int scalar.
  */
 int mlx_array_set_int(mlx_array* arr, int val);
+/**
+ * Set array to a float32 scalar.
+ */
+int mlx_array_set_float32(mlx_array* arr, float val);
 /**
  * Set array to a float scalar.
  */
 int mlx_array_set_float(mlx_array* arr, float val);
+/**
+ * Set array to a float64 scalar.
+ */
+int mlx_array_set_float64(mlx_array* arr, double val);
+/**
+ * Set array to a double scalar.
+ */
+int mlx_array_set_double(mlx_array* arr, double val);
 /**
  * Set array to a complex scalar.
  */
@@ -167,6 +194,7 @@ int mlx_array_dim(const mlx_array arr, int dim);
  * The array element type.
  */
 mlx_dtype mlx_array_dtype(const mlx_array arr);
+
 /**
  * Evaluate the array.
  */
@@ -212,6 +240,10 @@ int mlx_array_item_int64(int64_t* res, const mlx_array arr);
  * Access the value of a scalar array.
  */
 int mlx_array_item_float32(float* res, const mlx_array arr);
+/**
+ * Access the value of a scalar array.
+ */
+int mlx_array_item_float64(double* res, const mlx_array arr);
 /**
  * Access the value of a scalar array.
  */
@@ -281,6 +313,11 @@ const int64_t* mlx_array_data_int64(const mlx_array arr);
  * Array must be evaluated, otherwise returns NULL.
  */
 const float* mlx_array_data_float32(const mlx_array arr);
+/**
+ * Returns a pointer to the array data, cast to `float64*`.
+ * Array must be evaluated, otherwise returns NULL.
+ */
+const double* mlx_array_data_float64(const mlx_array arr);
 /**
  * Returns a pointer to the array data, cast to `_Complex*`.
  * Array must be evaluated, otherwise returns NULL.
@@ -302,6 +339,37 @@ const float16_t* mlx_array_data_float16(const mlx_array arr);
  */
 const bfloat16_t* mlx_array_data_bfloat16(const mlx_array arr);
 #endif
+
+/**
+ * Check if the array is available.
+ * Internal function: use at your own risk.
+ */
+int _mlx_array_is_available(bool* res, const mlx_array arr);
+
+/**
+ * Wait on the array to be available. After this `_mlx_array_is_available`
+ * returns `true`. Internal function: use at your own risk.
+ */
+int _mlx_array_wait(const mlx_array arr);
+
+/**
+ * Whether the array is contiguous in memory.
+ * Internal function: use at your own risk.
+ */
+int _mlx_array_is_contiguous(bool* res, const mlx_array arr);
+
+/**
+ * Whether the array's rows are contiguous in memory.
+ * Internal function: use at your own risk.
+ */
+int _mlx_array_is_row_contiguous(bool* res, const mlx_array arr);
+
+/**
+ * Whether the array's columns are contiguous in memory.
+ * Internal function: use at your own risk.
+ */
+int _mlx_array_is_col_contiguous(bool* res, const mlx_array arr);
+
 /**@}*/
 
 #ifdef __cplusplus

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

@@ -31,6 +31,7 @@ typedef enum mlx_compile_mode_ {
   MLX_COMPILE_MODE_NO_FUSE,
   MLX_COMPILE_MODE_ENABLED
 } mlx_compile_mode;
+int mlx_compile(mlx_closure* res, const mlx_closure fun, bool shapeless);
 int mlx_detail_compile(
     mlx_closure* res,
     const mlx_closure fun,

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

@@ -48,10 +48,18 @@ int mlx_device_set(mlx_device* dev, const mlx_device src);
  * Get device description.
  */
 int mlx_device_tostring(mlx_string* str, mlx_device dev);
+/**
+ * Check if devices are the same.
+ */
+bool mlx_device_equal(mlx_device lhs, mlx_device rhs);
+/**
+ * Returns the index of the device.
+ */
+int mlx_device_get_index(int* index, mlx_device dev);
 /**
  * Returns the type of the device.
  */
-mlx_device_type mlx_device_get_type(mlx_device dev);
+int mlx_device_get_type(mlx_device_type* type, mlx_device dev);
 /**
  * Returns the default MLX device.
  */

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

@@ -53,6 +53,12 @@ int mlx_linalg_eigvalsh(
     const char* UPLO,
     const mlx_stream s);
 int mlx_linalg_inv(mlx_array* res, const mlx_array a, const mlx_stream s);
+int mlx_linalg_lu(mlx_vector_array* res, const mlx_array a, const mlx_stream s);
+int mlx_linalg_lu_factor(
+    mlx_array* res_0,
+    mlx_array* res_1,
+    const mlx_array a,
+    const mlx_stream s);
 int mlx_linalg_norm_p(
     mlx_array* res,
     const mlx_array a,
@@ -82,6 +88,17 @@ int mlx_linalg_qr(
     mlx_array* res_1,
     const mlx_array a,
     const mlx_stream s);
+int mlx_linalg_solve(
+    mlx_array* res,
+    const mlx_array a,
+    const mlx_array b,
+    const mlx_stream s);
+int mlx_linalg_solve_triangular(
+    mlx_array* res,
+    const mlx_array a,
+    const mlx_array b,
+    bool upper,
+    const mlx_stream s);
 int mlx_linalg_svd(
     mlx_vector_array* res,
     const mlx_array a,

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

@@ -145,7 +145,7 @@ int mlx_as_strided(
     const mlx_array a,
     const int* shape,
     size_t shape_num,
-    const size_t* strides,
+    const int64_t* strides,
     size_t strides_num,
     size_t offset,
     const mlx_stream s);
@@ -162,6 +162,7 @@ int mlx_bitwise_and(
     const mlx_array a,
     const mlx_array b,
     const mlx_stream s);
+int mlx_bitwise_invert(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_bitwise_or(
     mlx_array* res,
     const mlx_array a,
@@ -473,6 +474,11 @@ int mlx_isinf(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_isnan(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_isneginf(mlx_array* res, const mlx_array a, const mlx_stream s);
 int mlx_isposinf(mlx_array* res, const mlx_array a, const mlx_stream s);
+int mlx_kron(
+    mlx_array* res,
+    const mlx_array a,
+    const mlx_array b,
+    const mlx_stream s);
 int mlx_left_shift(
     mlx_array* res,
     const mlx_array a,
@@ -757,6 +763,13 @@ int mlx_scatter_add(
     const int* axes,
     size_t axes_num,
     const mlx_stream s);
+int mlx_scatter_add_axis(
+    mlx_array* res,
+    const mlx_array a,
+    const mlx_array indices,
+    const mlx_array values,
+    int axis,
+    const mlx_stream s);
 int mlx_scatter_max(
     mlx_array* res,
     const mlx_array a,
@@ -960,6 +973,13 @@ 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_unflatten(
+    mlx_array* res,
+    const mlx_array a,
+    int axis,
+    const int* shape,
+    size_t shape_num,
+    const mlx_stream s);
 int mlx_var(
     mlx_array* res,
     const mlx_array a,

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

@@ -101,7 +101,7 @@ int mlx_random_permutation(
     int axis,
     const mlx_array key /* may be null */,
     const mlx_stream s);
-int mlx_random_permutation_all(
+int mlx_random_permutation_arange(
     mlx_array* res,
     int x,
     const mlx_array key /* may be null */,
@@ -116,7 +116,7 @@ int mlx_random_randint(
     const mlx_array key /* may be null */,
     const mlx_stream s);
 int mlx_random_seed(uint64_t seed);
-int mlx_random_split_equal_parts(
+int mlx_random_split_num(
     mlx_array* res,
     const mlx_array key,
     int num,

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

@@ -53,6 +53,10 @@ bool mlx_stream_equal(mlx_stream lhs, mlx_stream rhs);
  * Return the device of the stream.
  */
 int mlx_stream_get_device(mlx_device* dev, mlx_stream stream);
+/**
+ * Return the index of the stream.
+ */
+int mlx_stream_get_index(int* index, mlx_stream stream);
 /**
  * Synchronize with the provided stream.
  */