Repository navigation
idea for doing synchronous evals but _not_ holding evalLock - #474
davidkoski wants to merge 6 commits into
Conversation
| mlx_eval(vector_array) | ||
| mlx_async_eval(vector_array) | ||
| } | ||
| mlx_eval(vector_array) |
There was a problem hiding this comment.
This is the heart of the change. we need the evalLock to guard against mutation of shared state. python uses thread-local streams to scope this, but we can't do that on the swift side (at least not in the same way).
But what if we did the mutation part under lock and then just wait for the arrays to be ready outside the lock? The MLXArray instances are still not thread safe, so they can't escape to other threads where the mutation could be observed, so I think this is safe.
There was a problem hiding this comment.
FYI @aleroot -- I got this idea from the concurrent read work you did and this might be generally useful.
|
|
||
| @testable import MLX | ||
|
|
||
| // MARK: - Support |
There was a problem hiding this comment.
For now this is just some tests that I had claude throw together.
598ea73 to
ff38724
Compare
| ```swift | ||
| let g = grad { (x: MLXArray) -> MLXArray in | ||
| let y = x * x | ||
| print(y.asArray(Float.self)) // OK: synchronous eval of an intermediate |
There was a problem hiding this comment.
Clarify that this case is explicitly allowed -- not all transforms prohibit evaluation:
| @Suite("eval inside function transformations") | ||
| struct EvalInTransformsTests { |
There was a problem hiding this comment.
Based on PR feedback this tests and asserts the different evaluations inside of transforms that are allowed vs not.
| mlx_export_function_kwargs(url.path, closure, positionalArgs, kwargs, shapeless) | ||
| withEvalLock { | ||
| mlx_export_function_kwargs(url.path, closure, positionalArgs, kwargs, shapeless) |
There was a problem hiding this comment.
While investigating #499 found these spots that also require the evalLock
| let waitOutsideLock = withEvalLockState { state in | ||
| switch state { | ||
| case .nested: | ||
| // cannot release the lock to wait, and async eval rejects | ||
| // vjp/jvp tracers | ||
| _ = mlx_eval(vectorArray) | ||
| return false | ||
|
|
||
| case .outermost: | ||
| // schedule under the lock, wait (below) outside it | ||
| return mlx_async_eval(vectorArray) == 0 | ||
| } | ||
| } |
There was a problem hiding this comment.
This handling fixes the issues found in review, see also EvalInTransformsTests
- do the synchronous mutation under lock with asyncEval - synchronous eval then does not need lock
ff38724 to
4d773ae
Compare
This branch has not been deployed
3 participants
Proposed changes
Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes