mlx-bun 0.4.0 → 1.0.0

package/src/mlx/ffi.ts

@@ -1,512 +0,0 @@
-// Raw bun:ffi bindings to libmlxc.
-//
-// Calling convention (verified in Phase 0, see PLAN.md):
-// - Every mlx-c handle is a one-pointer struct `{ void* ctx }` — passed and
-//   returned by value in a single register on arm64, so FFIType.u64 works.
-// - Out-params (`mlx_array* res`) are a pointer to that slot: allocate a
-//   BigUint64Array(1), pass ptr(), reread after the call via read.u64 —
-//   NOT slot[0] (DFG stale-read bug, see rule above outArray).
-// - Ops return int status; 0 = success.
-// - mlx_optional_int is `{ int32 value; bool has_value }` = 8 bytes, by
-//   value in one register: value in bits 0..31, has_value at bit 32.
-
-import { dlopen, FFIType, ptr, read, type Pointer } from "bun:ffi";
-import { existsSync } from "node:fs";
-import { dirname, join } from "node:path";
-import { MLX_CORE_VERSION, nativePackDir } from "../native-pack";
-import { runtimeValue } from "../runtime-config";
-
-const { ptr: P, i32, u64, f32, f64, cstring } = FFIType;
-
-/** libmlxc resolution, in order: explicit env override → next to the
- *  executable (the `bun build --compile` sidecar layout, where
- *  libmlxc.dylib + libmlx.dylib + mlx.metallib ship beside the binary —
- *  see docs/reference/distribution.md) → the downloaded native-pack cache
- *  (src/native-pack.ts; populated on first run by the CLI) → homebrew
- *  (arm64, then Intel prefix). Keep in sync with nativeRuntimeDir(). */
-function resolveLibmlxc(): string {
-  const env = runtimeValue("MLX_BUN_LIBMLXC");
-  if (env) return env;
-  const beside = join(dirname(process.execPath), "libmlxc.dylib");
-  if (existsSync(beside)) return beside;
-  const cached = join(nativePackDir(), "libmlxc.dylib");
-  if (existsSync(cached)) return cached;
-  for (const p of ["/opt/homebrew/lib/libmlxc.dylib", "/usr/local/lib/libmlxc.dylib"])
-    if (existsSync(p)) return p;
-  // let dlopen produce the canonical error message for the default path
-  return "/opt/homebrew/lib/libmlxc.dylib";
-}
-
-export const LIBMLXC_PATH = resolveLibmlxc();
-export const C = dlopen(LIBMLXC_PATH, {
-  // linked MLX core version, independent of the C API family
-  mlx_version: { args: [P], returns: i32 },
-  mlx_string_new: { args: [], returns: u64 },
-  mlx_string_data: { args: [u64], returns: cstring },
-  mlx_string_free: { args: [u64], returns: i32 },
-  // array lifecycle
-  mlx_array_new: { args: [], returns: u64 },
-  mlx_array_new_data: { args: [P, P, i32, i32], returns: u64 },
-  mlx_array_new_data_managed_payload: { args: [P, P, i32, i32, P, P], returns: u64 },
-  mlx_array_free: { args: [u64], returns: i32 },
-  mlx_array_eval: { args: [u64], returns: i32 },
-  // array metadata
-  mlx_array_ndim: { args: [u64], returns: u64 },
-  mlx_array_shape: { args: [u64], returns: P },
-  mlx_array_dim: { args: [u64, i32], returns: i32 },
-  mlx_array_dtype: { args: [u64], returns: i32 },
-  mlx_array_size: { args: [u64], returns: u64 },
-  mlx_array_nbytes: { args: [u64], returns: u64 },
-  // data access (forces eval/copy to host-visible memory)
-  mlx_array_data_float32: { args: [u64], returns: P },
-  mlx_array_data_float16: { args: [u64], returns: P },
-  mlx_array_data_bfloat16: { args: [u64], returns: P },
-  mlx_array_data_uint32: { args: [u64], returns: P },
-  mlx_array_data_int32: { args: [u64], returns: P },
-  mlx_array_data_uint8: { args: [u64], returns: P },
-  mlx_array_data_int8: { args: [u64], returns: P },
-  mlx_array_item_float32: { args: [P, u64], returns: i32 },
-  // streams
-  mlx_default_gpu_stream_new: { args: [], returns: u64 },
-  mlx_default_cpu_stream_new: { args: [], returns: u64 },
-  mlx_stream_free: { args: [u64], returns: i32 },
-  // eval control (both take mlx_vector_array)
-  mlx_eval: { args: [u64], returns: i32 },
-  mlx_async_eval: { args: [u64], returns: i32 },
-  // memory introspection / limits
-  mlx_get_active_memory: { args: [P], returns: i32 },
-  mlx_get_peak_memory: { args: [P], returns: i32 },
-  mlx_reset_peak_memory: { args: [], returns: i32 },
-  // mlx_device_info is a one-pointer struct like every other handle
-  mlx_device_info_new: { args: [], returns: u64 },
-  mlx_device_info_get: { args: [P, u64], returns: i32 },
-  mlx_device_info_free: { args: [u64], returns: i32 },
-  mlx_device_info_get_size: { args: [P, u64, cstring], returns: i32 },
-  mlx_device_info_get_string: { args: [P, u64, cstring], returns: i32 },
-  mlx_get_default_device: { args: [P], returns: i32 },
-  mlx_device_free: { args: [u64], returns: i32 },
-  mlx_synchronize: { args: [u64], returns: i32 },
-  mlx_get_cache_memory: { args: [P], returns: i32 },
-  mlx_clear_cache: { args: [], returns: i32 },
-  mlx_set_memory_limit: { args: [P, u64], returns: i32 },
-  mlx_set_wired_limit: { args: [P, u64], returns: i32 },
-  // io: native lazy safetensors loader (Load primitive is CPU-only —
-  // always pass the CPU stream; GPU ops consume the unified buffers)
-  mlx_load_safetensors: { args: [P, P, cstring, u64], returns: i32 },
-  mlx_map_string_to_array_new: { args: [], returns: u64 },
-  mlx_map_string_to_array_free: { args: [u64], returns: i32 },
-  mlx_map_string_to_array_get: { args: [P, u64, cstring], returns: i32 },
-  mlx_map_string_to_string_new: { args: [], returns: u64 },
-  mlx_map_string_to_string_free: { args: [u64], returns: i32 },
-  // fast fused ops
-  // (res, x, weight may-be-null, eps, stream)
-  mlx_fast_rms_norm: { args: [P, u64, u64, f32, u64], returns: i32 },
-  // (res, x, weight may-be-null, bias may-be-null, eps, stream)
-  mlx_fast_layer_norm: { args: [P, u64, u64, u64, f32, u64], returns: i32 },
-  // MLX 0.32.2: mask, sinks, force_fused, stream.
-  mlx_fast_scaled_dot_product_attention: {
-    args: [P, u64, u64, u64, f32, cstring, u64, u64, FFIType.bool, u64], returns: i32,
-  },
-  mlx_matmul: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_contiguous: { args: [P, u64, FFIType.bool, u64], returns: i32 },
-  mlx_copy: { args: [P, u64, u64], returns: i32 },
-  // (res: mlx_vector_array*, w, opt group_size, opt bits, mode, global_scale, stream)
-  mlx_quantize: { args: [P, u64, u64, u64, cstring, u64, u64], returns: i32 },
-  mlx_vector_array_get: { args: [P, u64, u64], returns: i32 },
-  mlx_vector_array_new: { args: [], returns: u64 },
-  mlx_softmax_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_max_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_sum_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  // Optional dtype uses the eight-byte value/has_value layout of optional int.
-  mlx_cumsum_axis: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64, u64], returns: i32 },
-  mlx_cummax_axis: { args: [P, u64, i32, FFIType.bool, FFIType.bool, u64], returns: i32 },
-  mlx_maximum: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_expand_dims: { args: [P, u64, i32, u64], returns: i32 },
-  mlx_logical_or: { args: [P, u64, u64, u64], returns: i32 },
-  // (res, x, dims, traditional, opt base, scale, offset, freqs may-be-null, stream)
-  mlx_fast_rope: { args: [P, u64, i32, FFIType.bool, u64, f32, i32, u64, u64], returns: i32 },
-  // ops (grown as phases need them)
-  mlx_add: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_subtract: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_multiply: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_divide: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_tanh: { args: [P, u64, u64], returns: i32 },
-  mlx_sigmoid: { args: [P, u64, u64], returns: i32 },
-  mlx_erf: { args: [P, u64, u64], returns: i32 },
-  // (res, c, a, b, alpha, beta, stream) — alpha·(a@b) + beta·c (nn.Linear bias path)
-  mlx_addmm: { args: [P, u64, u64, u64, f32, f32, u64], returns: i32 },
-  mlx_power: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_take_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
-  mlx_reshape: { args: [P, u64, P, u64, u64], returns: i32 },
-  mlx_transpose_axes: { args: [P, u64, P, u64, u64], returns: i32 },
-  mlx_argmax_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_concatenate_axis: { args: [P, u64, i32, u64], returns: i32 },
-  mlx_vector_array_new_data: { args: [P, u64], returns: u64 },
-  mlx_vector_array_free: { args: [u64], returns: i32 },
-  mlx_array_item_uint32: { args: [P, u64], returns: i32 },
-  mlx_array_item_int32: { args: [P, u64], returns: i32 },
-  mlx_slice: { args: [P, u64, P, u64, P, u64, P, u64, u64], returns: i32 },
-  // (res, src, update, start*, n, stop*, n, strides*, n, stream)
-  mlx_slice_update: { args: [P, u64, u64, P, u64, P, u64, P, u64, u64], returns: i32 },
-  mlx_zeros: { args: [P, P, u64, i32, u64], returns: i32 },
-  mlx_exp: { args: [P, u64, u64], returns: i32 },
-  mlx_cos: { args: [P, u64, u64], returns: i32 },
-  mlx_sin: { args: [P, u64, u64], returns: i32 },
-  mlx_negative: { args: [P, u64, u64], returns: i32 },
-  mlx_argsort_axis: { args: [P, u64, i32, u64], returns: i32 },
-  mlx_argpartition_axis: { args: [P, u64, i32, i32, u64], returns: i32 },
-  mlx_arange: { args: [P, f64, f64, f64, i32, u64], returns: i32 },
-  mlx_take_along_axis: { args: [P, u64, u64, i32, u64], returns: i32 },
-  mlx_put_along_axis: { args: [P, u64, u64, u64, i32, u64], returns: i32 },
-  mlx_where: { args: [P, u64, u64, u64, u64], returns: i32 },
-  mlx_logsumexp_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_greater_equal: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_less: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_less_equal: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_logical_and: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_logical_not: { args: [P, u64, u64], returns: i32 },
-  mlx_equal: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_all: { args: [P, u64, FFIType.bool, u64], returns: i32 },
-  mlx_any_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_random_key: { args: [P, u64], returns: i32 },
-  mlx_random_categorical: { args: [P, u64, i32, u64, u64], returns: i32 },
-  mlx_astype: { args: [P, u64, i32, u64], returns: i32 },
-  // (res, w, scales, biases, opt group_size, opt bits, mode, global_scale, opt dtype, stream)
-  mlx_dequantize: { args: [P, u64, u64, u64, u64, u64, cstring, u64, u64, u64], returns: i32 },
-  mlx_quantized_matmul: { args: [P, u64, u64, u64, u64, FFIType.bool, u64, u64, cstring, u64], returns: i32 },
-  // (res, x, w, scales, biases?, lhs_indices?, rhs_indices?, transpose,
-  //  opt group_size, opt bits, mode, sorted_indices, stream)
-  mlx_gather_qmm: { args: [P, u64, u64, u64, u64, u64, u64, FFIType.bool, u64, u64, cstring, FFIType.bool, u64], returns: i32 },
-  mlx_floor_divide: { args: [P, u64, u64, u64], returns: i32 },
-  // (res, input, weight, stride, padding, dilation, groups, stream) — depthwise
-  // causal conv1d for Qwen3.5 gated-DeltaNet (groups == channels, padding 0).
-  mlx_conv1d: { args: [P, u64, u64, i32, i32, i32, i32, u64], returns: i32 },
-  // Natural mlx-c signatures. Bun 1.4 fixed the macOS arm64 stack packing
-  // bug recorded in lab/repro/bun-ffi-stack-args, so these must match the
-  // header exactly. mlx-bun now requires Bun 1.4 or newer.
-  mlx_conv2d: { args: [P, u64, u64, i32, i32, i32, i32, i32, i32, i32, u64], returns: i32 },
-  mlx_conv3d: { args: [P, u64, u64, i32, i32, i32, i32, i32, i32, i32, i32, i32, i32, u64], returns: i32 },
-  // --- training: autograd (value_and_grad) — proven in lab/spikes/phase-train-vag.ts.
-  // mlx_value_and_grad(res: mlx_closure_value_and_grad*, fun: mlx_closure,
-  //   argnums: int*, n) → builds a value+grad closure differentiating the
-  //   selected inputs of `fun` (a mlx_closure_new_func_payload over a JS loss).
-  mlx_value_and_grad: { args: [P, u64, P, u64], returns: i32 },
-  mlx_closure_value_and_grad_new: { args: [], returns: u64 },
-  // (value: vector_array*, grads: vector_array*, cls, input: vector_array)
-  mlx_closure_value_and_grad_apply: { args: [P, P, u64, u64], returns: i32 },
-  mlx_closure_value_and_grad_free: { args: [u64], returns: i32 },
-  // mlx_vjp(outputs: vector_array*, vjps: vector_array*, fun: mlx_closure,
-  //   primals: vector_array, cotangents: vector_array) — reverse-mode AD given an
-  //   explicit cotangent per output. The natural segmented-backward primitive:
-  //   fun = segment forward, cotangent = dh, vjps = [dh_in, ...dLoRA]. No
-  //   surrogate scalar, no value_and_grad closure (see src/train/segmented.ts).
-  mlx_vjp: { args: [P, P, u64, u64, u64], returns: i32 },
-  // --- training: random init (LoRA A ~ normal; B stays zeros)
-  // (res, shape*, n, dtype, loc, scale, key may-be-null, stream)
-  mlx_random_normal: { args: [P, P, u64, i32, f32, f32, u64, u64], returns: i32 },
-  // (res, low, high, shape*, n, dtype, key may-be-null, stream) — low/high are
-  // scalar mlx_arrays. Used for mlx-lm-faithful LoRA A init (uniform).
-  mlx_random_uniform: { args: [P, u64, u64, P, u64, i32, u64, u64], returns: i32 },
-  mlx_random_randint: { args: [P, u64, u64, P, u64, i32, u64, u64], returns: i32 },
-  mlx_random_seed: { args: [u64], returns: i32 },
-  mlx_random_split_num: { args: [P, u64, i32, u64], returns: i32 },
-  // --- training: loss/optimizer elementwise + reductions
-  mlx_stop_gradient: { args: [P, u64, u64], returns: i32 },
-  mlx_log: { args: [P, u64, u64], returns: i32 },
-  mlx_logaddexp: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_mean: { args: [P, u64, FFIType.bool, u64], returns: i32 },
-  mlx_mean_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_sqrt: { args: [P, u64, u64], returns: i32 },
-  mlx_rsqrt: { args: [P, u64, u64], returns: i32 },
-  mlx_square: { args: [P, u64, u64], returns: i32 },
-  mlx_abs: { args: [P, u64, u64], returns: i32 },
-  // (res, a, a_min may-be-null, a_max may-be-null, stream)
-  mlx_clip: { args: [P, u64, u64, u64, u64], returns: i32 },
-  // --- quantization: native safetensors writer (io.h) + map inserts (map.h)
-  mlx_save_safetensors: { args: [cstring, u64, u64], returns: i32 },
-  mlx_map_string_to_array_insert: { args: [u64, cstring, u64], returns: i32 },
-  mlx_map_string_to_string_insert: { args: [u64, cstring, cstring], returns: i32 },
-  // compile via closures (Phase A compiled decode — docs/archive/investigations/optimization_plan.md).
-  // mlx_closure is the usual one-pointer struct; the func-payload variant
-  // carries an id we use to route the trace callback back to JS.
-  mlx_closure_new: { args: [], returns: u64 },
-  mlx_closure_new_func_payload: { args: [P, P, P], returns: u64 },
-  mlx_closure_free: { args: [u64], returns: i32 },
-  mlx_closure_apply: { args: [P, u64, u64], returns: i32 },
-  // gradient checkpointing (transforms.h): wrap a closure so its forward
-  // activations are dropped and recomputed during the enclosing backward.
-  // mlx_checkpoint(res: mlx_closure*, fun: mlx_closure) → checkpointed closure.
-  mlx_checkpoint: { args: [P, u64], returns: i32 },
-  // custom_vjp (transforms.h): attach a hand-written backward to a forward whose
-  // primitives have no usable vjp (e.g. a fused metal_kernel). Needed for the L2
-  // flash-attention training op. mlx_custom_vjp(res: mlx_closure*, fun: mlx_closure,
-  //   fun_vjp: mlx_closure_custom). The custom-vjp closure's callback is
-  //   (res*, primals, cotangents, outputs, payload) → status.
-  mlx_custom_vjp: { args: [P, u64, u64], returns: i32 },
-  mlx_closure_custom_new_func_payload: { args: [P, P, P], returns: u64 },
-  mlx_closure_custom_free: { args: [u64], returns: i32 },
-  mlx_compile: { args: [P, u64, FFIType.bool], returns: i32 },
-  mlx_set_compile_mode: { args: [i32], returns: i32 },
-  mlx_vector_array_size: { args: [u64], returns: u64 },
-  mlx_vector_array_set_data: { args: [P, P, u64], returns: i32 },
-  // dynamic-start variants: per-decode-step offsets enter the graph as
-  // array VALUES instead of baked ints, so one compiled graph serves
-  // every step (no per-step retrace).
-  // (res, x, dims, traditional, opt base, scale, offset ARRAY, freqs may-be-null, stream)
-  mlx_fast_rope_dynamic: { args: [P, u64, i32, FFIType.bool, u64, f32, u64, u64, u64], returns: i32 },
-  // (res, src, update, start ARRAY, axes*, axes_num, stream)
-  mlx_slice_update_dynamic: { args: [P, u64, u64, u64, P, u64, u64], returns: i32 },
-  // (res, a, start ARRAY, axes*, axes_num, slice_size*, n, stream)
-  mlx_slice_dynamic: { args: [P, u64, u64, P, u64, P, u64, u64], returns: i32 },
-  // custom Metal kernels (Phase E — mx.fast.metal_kernel from Bun)
-  mlx_vector_string_new_data: { args: [P, u64], returns: u64 },
-  mlx_vector_string_free: { args: [u64], returns: i32 },
-  // (name, input_names, output_names, source, header, ensure_row_contiguous, atomic_outputs)
-  mlx_fast_metal_kernel_new: { args: [cstring, u64, u64, cstring, cstring, FFIType.bool, FFIType.bool], returns: u64 },
-  mlx_fast_metal_kernel_free: { args: [u64], returns: FFIType.void },
-  // (outputs: mlx_vector_array*, kernel, inputs, config, stream)
-  mlx_fast_metal_kernel_apply: { args: [P, u64, u64, u64, u64], returns: i32 },
-  mlx_fast_metal_kernel_config_new: { args: [], returns: u64 },
-  mlx_fast_metal_kernel_config_free: { args: [u64], returns: FFIType.void },
-  mlx_fast_metal_kernel_config_add_output_arg: { args: [u64, P, u64, i32], returns: i32 },
-  mlx_fast_metal_kernel_config_set_grid: { args: [u64, i32, i32, i32], returns: i32 },
-  mlx_fast_metal_kernel_config_set_thread_group: { args: [u64, i32, i32, i32], returns: i32 },
-  mlx_fast_metal_kernel_config_set_init_value: { args: [u64, f32], returns: i32 },
-  mlx_fast_metal_kernel_config_add_template_arg_int: { args: [u64, cstring, i32], returns: i32 },
-  mlx_fast_metal_kernel_config_add_template_arg_dtype: { args: [u64, cstring, i32], returns: i32 },
-  // GPU trace capture (Phase E step 2: size the prize before the kernel)
-  mlx_metal_start_capture: { args: [cstring], returns: i32 },
-  mlx_metal_stop_capture: { args: [], returns: i32 },
-  // TurboQuant KV codec (Phase 13): FWHT rotation + per-group affine/Lloyd-Max
-  // (res, a, opt scale, stream) — 1/sqrt(dim) normalization when scale is null.
-  mlx_hadamard_transform: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_min_axis: { args: [P, u64, i32, FFIType.bool, u64], returns: i32 },
-  mlx_round: { args: [P, u64, i32, u64], returns: i32 },
-  mlx_greater: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_bitwise_and: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_bitwise_or: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_left_shift: { args: [P, u64, u64, u64], returns: i32 },
-  mlx_right_shift: { args: [P, u64, u64, u64], returns: i32 },
-}).symbols;
-
-export type MlxHandle = bigint;
-
-// --- error handling -------------------------------------------------------
-// mlx-c's default error handler prints and aborts the process. Install a
-// recording handler instead: ops then return non-zero status and outArray
-// throws a JS Error carrying the mlx message + JS stack. mlx-c invokes the
-// handler on the calling thread (exceptions are caught in the C wrappers),
-// so a non-threadsafe JSCallback is safe here.
-import { CString, JSCallback } from "bun:ffi";
-
-let lastMlxError: string | null = null;
-
-const errorHandler = new JSCallback(
-  (msgPtr: number) => {
-    lastMlxError = msgPtr ? new CString(msgPtr as never).toString() : "unknown mlx error";
-  },
-  { args: ["ptr", "ptr"], returns: "void" },
-);
-
-const errLib = dlopen(LIBMLXC_PATH, {
-  mlx_set_error_handler: { args: [P, P, P], returns: FFIType.void },
-}).symbols;
-errLib.mlx_set_error_handler(errorHandler.ptr, null, null);
-
-/** Consume the last recorded mlx error message (if any). */
-export function takeMlxError(): string | null {
-  const e = lastMlxError;
-  lastMlxError = null;
-  return e;
-}
-
-function readRuntimeVersion(): string {
-  const slot = new BigUint64Array([C.mlx_string_new()]);
-  const p = ptr(slot);
-  try {
-    if (C.mlx_version(p) !== 0)
-      throw new Error(`mlx_version failed: ${takeMlxError() ?? "(no mlx message)"}`);
-    return C.mlx_string_data(read.u64(p, 0)).toString();
-  } finally { C.mlx_string_free(read.u64(p, 0)); }
-}
-
-/** Read once at startup; kernel eligibility must follow the linked core. */
-export const MLX_VERSION = readRuntimeVersion();
-if (MLX_VERSION !== MLX_CORE_VERSION) {
-  throw new Error(`mlx-bun requires its bundled MLX ${MLX_CORE_VERSION}; ${LIBMLXC_PATH} loaded MLX ${MLX_VERSION}`);
-}
-
-/** mlx_dtype enum values (array.h). */
-export const enum Dtype {
-  bool = 0, uint8 = 1, uint16 = 2, uint32 = 3, uint64 = 4,
-  int8 = 5, int16 = 6, int32 = 7, int64 = 8,
-  float16 = 9, float32 = 10, float64 = 11, bfloat16 = 12, complex64 = 13,
-}
-
-export const DTYPE_NAMES: Record<number, string> = {
-  0: "bool", 1: "uint8", 2: "uint16", 3: "uint32", 4: "uint64",
-  5: "int8", 6: "int16", 7: "int32", 8: "int64",
-  9: "float16", 10: "float32", 11: "float64", 12: "bfloat16", 13: "complex64",
-};
-
-/** Pack mlx_optional_int for by-value passing. */
-export function optInt(value: number | null): bigint {
-  if (value === null) return 0n;
-  return (BigInt(value >>> 0) & 0xffffffffn) | (1n << 32n);
-}
-
-/** Pack mlx_optional_dtype (same layout: enum + has_value flag). */
-export const optDtype = optInt;
-
-/** Pack mlx_optional_float `{float value; bool has_value}` for by-value passing. */
-const f32Bits = new DataView(new ArrayBuffer(4));
-export function optFloat(value: number | null): bigint {
-  if (value === null) return 0n;
-  f32Bits.setFloat32(0, value, true);
-  return BigInt(f32Bits.getUint32(0, true)) | (1n << 32n);
-}
-
-/** Null handle for optional mlx_array parameters. */
-export const NULL_HANDLE = 0n;
-
-// RULE: never read a typed array that a bun:ffi call wrote through a
-// pointer — once the calling function is DFG-compiled, the JIT eliminates
-// the load across the native call and `buf[i]` returns stale values
-// (lab/repro/bun-ffi-f64/ISSUE.md; PLAN.md Phase 4 findings). Out-param slots
-// must be read back with bun:ffi `read.*` (verified safe). Initializing the
-// slot via the typed-array *constructor* is fine: that store happens in
-// host code the JIT can't elide.
-
-/** Run an op that writes its result through an mlx_array* out-param. */
-export function outArray(op: string, call: (slotPtr: number) => number): MlxHandle {
-  const slot = new BigUint64Array([C.mlx_array_new()]);
-  const slotPtr = ptr(slot);
-  const status = call(slotPtr);
-  const handle = read.u64(slotPtr, 0);
-  if (status !== 0) {
-    C.mlx_array_free(handle);
-    throw new Error(`${op} failed: ${takeMlxError() ?? "(no mlx message)"}`);
-  }
-  return handle;
-}
-
-export function activeMemory(): number {
-  const out = new BigUint64Array(1);
-  const p = ptr(out);
-  C.mlx_get_active_memory(p);
-  return Number(read.u64(p, 0));
-}
-
-export function cacheMemory(): number {
-  const out = new BigUint64Array(1);
-  const p = ptr(out);
-  C.mlx_get_cache_memory(p);
-  return Number(read.u64(p, 0));
-}
-
-export function peakMemory(): number {
-  const out = new BigUint64Array(1);
-  const p = ptr(out);
-  C.mlx_get_peak_memory(p);
-  return Number(read.u64(p, 0));
-}
-
-/** mx.reset_peak_memory — zero the peak counter (e.g. after model load,
- *  so a subsequent peakMemory() reads generation-only peak). */
-export function resetPeakMemory(): void {
-  C.mlx_reset_peak_memory();
-}
-
-/** mx.clear_cache — release the allocator's cached buffers back to the
- *  OS. mlx-lm calls this after every prefill chunk (and every 256 decode
- *  steps); without it the first decode step after a long prefill pays a
- *  one-shot allocator-reclaim stall (~800 ms after an 8k prefill —
- *  measured, scripts/decode-split.ts; the root cause of the
- *  context-scaling decode gap). */
-export function clearCache(): void {
-  C.mlx_clear_cache();
-}
-
-let cachedDeviceArchitecture: string | undefined;
-
-/** Native architecture identifier, used to match MLX's hardware dispatch. */
-export function deviceArchitecture(): string {
-  if (cachedDeviceArchitecture !== undefined) return cachedDeviceArchitecture;
-  const devSlot = new BigUint64Array(1), devPtr = ptr(devSlot);
-  if (C.mlx_get_default_device(devPtr) !== 0)
-    throw new Error(`mlx_get_default_device failed: ${takeMlxError() ?? ""}`);
-  const dev = read.u64(devPtr, 0);
-  const infoSlot = new BigUint64Array([C.mlx_device_info_new()]), infoPtr = ptr(infoSlot);
-  try {
-    if (C.mlx_device_info_get(infoPtr, dev) !== 0)
-      throw new Error(`mlx_device_info_get failed: ${takeMlxError() ?? ""}`);
-    const out = new BigUint64Array(1), outPtr = ptr(out);
-    const key = Buffer.from("architecture\0", "utf8");
-    if (C.mlx_device_info_get_string(outPtr, read.u64(infoPtr, 0), ptr(key)) !== 0)
-      throw new Error("architecture not in device info");
-    // Copy before the device-info object releases the native string.
-    return cachedDeviceArchitecture = new CString(Number(read.u64(outPtr, 0)) as Pointer).toString();
-  } finally {
-    C.mlx_device_info_free(read.u64(infoPtr, 0));
-    C.mlx_device_free(dev);
-  }
-}
-
-/** Metal's recommended max working-set size for the default device. */
-export function maxRecommendedWorkingSetSize(): number {
-  const devSlot = new BigUint64Array(1);
-  const devPtr = ptr(devSlot);
-  if (C.mlx_get_default_device(devPtr) !== 0)
-    throw new Error(`mlx_get_default_device failed: ${takeMlxError() ?? ""}`);
-  const dev = read.u64(devPtr, 0);
-  const infoSlot = new BigUint64Array([C.mlx_device_info_new()]);
-  const infoPtr = ptr(infoSlot);
-  try {
-    if (C.mlx_device_info_get(infoPtr, dev) !== 0)
-      throw new Error(`mlx_device_info_get failed: ${takeMlxError() ?? ""}`);
-    const info = read.u64(infoPtr, 0);
-    const out = new BigUint64Array(1);
-    const outPtr = ptr(out);
-    const key = Buffer.from("max_recommended_working_set_size\0", "utf8");
-    if (C.mlx_device_info_get_size(outPtr, info, ptr(key)) !== 0)
-      throw new Error("max_recommended_working_set_size not in device info");
-    return Number(read.u64(outPtr, 0));
-  } finally {
-    C.mlx_device_info_free(read.u64(infoPtr, 0));
-    C.mlx_device_free(dev);
-  }
-}
-
-/** mx.set_wired_limit — returns the previous limit. Models close to the
- *  working-set ceiling decode ~4x slower without this (Metal evicts and
- *  re-faults weight buffers every token; mlx-lm's wired_limit context
- *  is the reference behavior). MUST be scoped (set → generate →
- *  synchronize → restore), never process-permanent: a permanent limit
- *  pins buffers across idle periods and OOM-kills multi-model processes
- *  (the test suite) via uncatchable async GPU errors. */
-export function setWiredLimit(bytes: number): number {
-  const out = new BigUint64Array(1);
-  const p = ptr(out);
-  if (C.mlx_set_wired_limit(p, BigInt(Math.floor(bytes))) !== 0)
-    throw new Error(`mlx_set_wired_limit failed: ${takeMlxError() ?? ""}`);
-  return Number(read.u64(p, 0));
-}
-
-/** mx.set_memory_limit — returns the previous limit. At the limit the
- *  allocator reclaims/waits instead of ballooning past it. Defense in
- *  depth under admission control ONLY: it does NOT make Metal
- *  command-buffer OOM catchable (Phase 6 finding — that throw comes
- *  from a completion handler and is std::terminate). */
-export function setMemoryLimit(bytes: number): number {
-  const out = new BigUint64Array(1);
-  const p = ptr(out);
-  if (C.mlx_set_memory_limit(p, BigInt(Math.floor(bytes))) !== 0)
-    throw new Error(`mlx_set_memory_limit failed: ${takeMlxError() ?? ""}`);
-  return Number(read.u64(p, 0));
-}
-
-/** mx.synchronize(stream) — wired limit must not change mid-async-eval. */
-export function synchronize(stream: MlxHandle): void {
-  if (C.mlx_synchronize(stream) !== 0)
-    throw new Error(`mlx_synchronize failed: ${takeMlxError() ?? ""}`);
-}

package/src/mlx/materialize.ts

@@ -1,23 +0,0 @@
-// Explicit compact storage for a value that must release a larger backing
-// buffer. MLX 0.31.2 Copy::eval shares storage; contiguous() may also alias.
-// MLX's dynamic slice allocates the logical output size and has a native VJP.
-import { MlxArray } from "./array";
-import * as ops from "./ops";
-
-let zero: MlxArray | undefined;
-
-/** Borrow an array; return an owned, row-contiguous copy on the GPU.
- * Values and dtype are unchanged. Evaluation stays lazy and never reads data
- * back to the host. Use only when independent compact storage is required;
- * immutable array aliases normally avoid this allocation and dispatch.
- * Slice sizes are concrete, so use shape-specialized compiled graphs. */
-export function materializeCopy(input: MlxArray): MlxArray {
-  zero ??= MlxArray.fromInt32(new Int32Array([0]), []);
-  if (input.ndim > 0) return ops.sliceDynamic(input, zero, [0], input.shape);
-  const vector = ops.reshape(input, [1]);
-  try {
-    const copy = ops.sliceDynamic(vector, zero, [0], [1]);
-    try { return ops.reshape(copy, []); }
-    finally { copy.dispose(); }
-  } finally { vector.dispose(); }
-}