Skip to content

Bind new_thread_unsafe_stream - #122

Open
CharlieTLe wants to merge 1 commit into
ml-explore:mainfrom
CharlieTLe:add-thread-unsafe-stream
Open

Bind new_thread_unsafe_stream#122
CharlieTLe wants to merge 1 commit into
ml-explore:mainfrom
CharlieTLe:add-thread-unsafe-stream

Conversation

@CharlieTLe

Copy link
Copy Markdown

Summary

Expose mlx's new_thread_unsafe_stream as mlx_stream_new_thread_unsafe.

Streams became thread affine in mlx 0.31.2 (#3348, #3281): a stream's GPU command encoder is registered in a thread_local map, so evaluating on a stream from a thread other than the one that created it raises There is no Stream(gpu, N) in current thread.

mlx added new_thread_unsafe_stream in ml-explore/mlx#3578 specifically for this case:

This provides a way for language bindings (like Swift) to evaluate graph asynchronously via threads, on the condition that they add locks on themselves.

mlx-c does not bind it, so those bindings can't reach it. mlx-swift currently carries a local C shim to work around exactly this, which this PR would let it delete.

Before / after

examples/example-thread-stream.c creates a stream on the main thread and evaluates on it from a worker thread. With mlx_stream_new_device:

MLX error: There is no Stream(gpu, 0) in current thread. at mlx/c/array.cpp:443

With mlx_stream_new_thread_unsafe:

2 + 3 evaluated on another thread = 5

Note on the mlx pin

new_thread_unsafe_stream only exists in mlx 0.32.0, so GIT_TAG moves from v0.31.2 to v0.32.0. I kept that to the one-line bump rather than regenerating bindings for 0.32.0, on the assumption you'd rather do that as its own change — happy to drop the bump if you'd prefer to land this after a 0.32.0 regen.

The existing bindings compile unchanged against 0.32.0, and example, example-grad, example-closure, example-graph and example-export all pass.

Naming

I went with mlx_stream_new_thread_unsafe to sit alongside mlx_stream_new_device in the mlx_stream_new* constructor family. Happy to rename to mlx_new_thread_unsafe_stream (closer to the C++ name) or anything else you prefer.

Possible follow-up

new_thread_local_stream / stream_from_thread_local_stream / clear_streams are also unbound. They need a new opaque mlx_thread_local_stream type with its own private accessors, so I left them out to keep this focused — glad to do it in a separate PR if it's wanted.

Checklist

  • Formatted with clang-format using the repo's .clang-format
  • Added an example exercising the new API
  • Built and ran locally (macOS 27.0, Apple Silicon)

Streams are thread affine as of mlx 0.31.2: a stream's GPU command encoder
is registered in a thread_local map, so evaluating on a stream from a thread
other than the one that created it raises "There is no Stream(gpu, N) in
current thread."

mlx added `new_thread_unsafe_stream` for language bindings that do their own
locking (ml-explore/mlx#3578), but mlx-c does not expose it, so those bindings
have no way to reach it. mlx-swift currently carries a local C shim for this.

Requires mlx 0.32.0, so the pin moves from v0.31.2. The existing bindings
build and the examples pass unchanged against it.
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