Skip to content

Add mlx_array_detach to break graph links before free#119

Open
odilitime wants to merge 1 commit into
ml-explore:mainfrom
odilitime:feat/array-detach
Open

Add mlx_array_detach to break graph links before free#119
odilitime wants to merge 1 commit into
ml-explore:mainfrom
odilitime:feat/array-detach

Conversation

@odilitime

Copy link
Copy Markdown

Summary

Adds mlx_array_detach() to the C API, wrapping the existing mlx::core::array::detach() method.

When releasing model weights or other arrays that may still be linked in the computation graph (e.g. sibling outputs from a traced/compiled forward pass), calling mlx_array_free() alone can leave dangling references in shared ArrayDesc nodes. Detaching first clears those sibling links so the buffer can be freed safely.

  • Declares mlx_array_detach in mlx/c/array.h with documentation
  • Implements a thin wrapper in mlx/c/array.cpp
  • No-op when arr.ctx is null, consistent with mlx_array_free_

Motivation

Downstream bindings (Go via dynamic libmlxc, Swift, etc.) that hold weight arrays across inference sessions currently need to patch array.cpp locally or call into C++ internals to detach before free. Exposing this in the public C API removes that friction.

Test plan

  • Build mlx-c against MLX v0.31.2
  • Verify mlx_array_detach is exported from libmlxc
  • Manual: create an array via ops, detach, free — no crash
  • Manual: detach + free weight arrays after a compiled forward pass — no use-after-free

Expose mlx::core::array::detach() in the C API so callers can release
weight buffers without leaving sibling references in shared ArrayDesc
nodes. No-op when arr.ctx is null, matching mlx_array_free behavior.
@odilitime

Copy link
Copy Markdown
Author

Test plan results (macOS arm64)

Verified locally on feat/array-detach (3da2b0a):

  • Build mlx-c against MLX v0.31.2cmake -B build-static -DCMAKE_BUILD_TYPE=Release fetched MLX v0.31.2; cmake --build build-static --target mlxc succeeded (static lib).
  • Verify mlx_array_detach exported from libmlxcnm build-static/libmlxc.a | rg mlx_array_detachT _mlx_array_detach.
  • Manual: ops → detach → free — add on GPU, detach operands + result + empty array, free; no crash.
  • Manual: compiled forward → detach weight → freemlx_compile + matmul closure using a weight array captured via payload; detach + free weight after apply; no crash/use-after-free.
MLX version: 0.31.2
PASS: create array via ops, detach, free
PASS: detach + free weight after compiled forward pass
All mlx_array_detach tests passed.

Note: BUILD_SHARED_LIBS=ON failed to link libjaccl.dylib on this host (missing C++ stdlib symbols); static libmlxc.a build and symbol export check above are unaffected.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant