-
Notifications
You must be signed in to change notification settings - Fork 37
Description
I'm building Rust bindings for MLX-C (mlx-rs) and encountered an ABI mismatch when calling functions that accept mlx_optional_int structs by value on ARM64 (Apple Silicon).
Functions like mlx_quantize, mlx_dequantize, and mlx_quantized_matmul take mlx_optional_int parameters by value. When called from Rust via FFI (with #[repr(C)] and matching struct sizes), the struct arguments corrupt subsequent parameters — specifically, the const char* mode parameter receives a garbage pointer instead of the intended string.
The same functions work correctly when called from C/C++. The issue is a subtle ABI lowering disagreement between rustc and clang on how small mixed-type structs ({int, bool}) are passed in registers on aarch64.
This affects any non-C/C++ language binding (Rust)that calls these functions via FFI. The bug is silent and difficult to diagnose — struct sizes match, #[repr(C)] is correct, but arguments silently shift at the calling convention level.
Provide alternative entry points that accept plain scalar types instead of mlx_optional_* structs:
int mlx_quantize_simple( mlx_vector_array* res, mlx_array w, int group_size, int bits, const char* mode, mlx_stream s);
These would be thin wrappers that construct the mlx_optional_* on the C side and forward to the existing implementation. This is the pattern I'm currently using as a workaround, and it resolves the issue completely. Alternatively, documenting this known pitfall would also be valuable.