Make the evaluation barrier a graph decision (Materialize modes) - #220
Conversation
Materialize forced a full device sync unconditionally. That is the strongest of three useful guarantees and the most expensive: mx.eval blocks the calling thread for a device round-trip, measured at ~0.1-0.15 ms per call on an M4 Pro regardless of array size. Often the guarantee actually wanted is weaker -- "the lazy graph must not accumulate across calls" -- and mx.async_eval provides exactly that, detaching the pending graph without the stall. Sometimes it is weaker still, because something downstream (a conversion back to NumPy, an outlet) already forces evaluation every cycle, and the right answer is to do nothing at all. MaterializeMode names all three so a graph can pick one per site, and materialize_array applies it to a single array so other nodes can share the policy rather than re-implementing the try/except ImportError dance. Default stays SYNC: this node exists only to be a barrier, so wiring one and getting something weaker would be surprising. Adds a regression test for the property the node exists for -- with OFF, a recurrence retains every input it has ever seen (52 MB over 100 iterations); with SYNC or ASYNC it retains one (0.53 MB). Documents the two things that decide where the barrier goes: evaluating an output also evaluates state emitted by the same multi-output kernel (so stateful transformers need no internal eval of their own), and evaluating an empty message forces nothing (so the barrier must sit upstream of any stage that can withhold output).
|
Worth flagging that this PR fixes a latent bug beyond the stated scope: On 3.2.0 it declares no Found it wiring This PR incidentally fixes it by giving the Unit |
Why
Materializeforced a full device sync unconditionally. That is the strongest of three useful guarantees and the most expensive:mx.evalblocks the calling thread for a device round-trip, measured at ~0.1–0.15 ms per call on an M4 Pro regardless of array size — it is pure latency, not work.Often the guarantee actually wanted is weaker: "the lazy graph must not accumulate across calls."
mx.async_evalprovides exactly that — it detaches the pending graph just asevaldoes, but schedules rather than blocks, leaving the CPU free to build the next message. Sometimes it is weaker still, because something downstream (a conversion back to NumPy, an outlet, a serializer) already forces evaluation every cycle, and the right answer is to do nothing.A lazy backend only pays off while the graph stays lazy, so where to force evaluation is a property of the graph, not of any one node. This makes that choice expressible.
What
MaterializeMode—sync/async/off, each documented with what it guarantees and what it costs.materialize_array(data, mode)— applies a mode to one array. A no-op on non-lazy backends. Exposed so other nodes can share the policy instead of re-implementing thetry/except ImportErrordance.MaterializeSettings/Materialize.SETTINGS/ amaterialize()factory, matching the convention inclip.pyetc.SYNC— this node exists only to be a barrier, so wiring one and getting something weaker would be surprising.Backward compatible:
MaterializeTransformer()and a bareMaterializeunit behave exactly as before.Test
New regression test for the property the node exists for. A recurrence that feeds its output back as state retains every input it has ever seen under
off(52 MB over 100 iterations) and exactly one undersyncorasync(0.53 MB).Full suite: 4063 passed, 6 skipped.
Docs
Adds a "Choosing where to evaluate" section covering the two facts that decide placement:
ewma_mlx_metalandsosfilt_mlx_metalemit filter state as a second output of the same kernel as the signal, so one barrier on the output covers both — which is why the stateful transformers need no internalmx.eval, and deliberately have none.