ezmsg.sigproc.materialize#
Materialize (evaluate) lazy array data.
MLX arrays are lazily evaluated — computations are queued but not executed until the result is needed. This module provides an explicit evaluation point so that downstream processors receive fully-evaluated data.
Placing the evaluation point is a graph decision, not a per-node one: a lazy
backend only pays off while the graph stays lazy, so every forced evaluation
trades throughput for a guarantee. MaterializeMode names the three
positions on that trade so a graph can pick one per site.
Where to place it#
mx.async_eval never blocks — not on a dependency still in flight (four
chained stages enqueued in 2.4 ms against 71 ms of device work), and not when
the queue is already deep (24 chunks enqueued in 5.7 ms against 201 ms, with no
stall). Which makes it tempting to evaluate after every node. Don’t: each
evaluation point costs ~17-21 µs of host time to submit a command buffer,
against ~1 µs to build the op graph it replaces — and in a live pipeline that
host time is the executor’s, not the device’s.
Measured on an M4 Pro across a three-node MLX segment (Metal SOS filter, scaler, abs), µs per message:
chunk |
SYNC at tail |
ASYNC tail |
every node |
|---|---|---|---|
30-64 x 256 ch |
300-334 |
139-172 |
201-223 |
128 x 1024 ch |
594-641 |
338-342 |
328-331 |
1024 x 2048 ch |
1961-2009 |
1602-1650 |
1502-1593 |
At streaming chunk sizes, evaluating per node is 1.3-1.7x worse than one evaluation at the tail. It only breaks even once a single node’s device work is large enough to bury the submission cost, which is well above any chunk a real-time graph carries. Fan-out is not a special case either: evaluating each branch as it finishes measured 1.3x worse than evaluating once at the join.
What does pay is putting that one evaluation point at the end of the MLX segment, before whatever else the executor has to do — downstream units, or the NumPy conversion at a process boundary (shared memory cannot carry MLX arrays, so that conversion forces completion regardless). Against deferring to that conversion, as a function of how much other work sits in between:
other host work |
deferred |
ASYNC at the end |
|---|---|---|
none |
1.00x |
1.00x |
~0.5 ms |
1.00x |
1.30-1.41x |
~2 ms |
1.00x |
1.63-1.69x |
~6 ms |
1.00x |
1.39-1.42x |
With nothing else to do the two are identical: the gain is entirely overlap
between device work and host work, so it is worth exactly as much as the host
has to get on with. benchmarks/benchmark_mlx_async_eval_placement.py
reproduces all of the above.
Note
Evaluating an empty message forces nothing. A branch that emits
zero-length outputs on some cycles (a binning stage fed sub-bin chunks, say)
leaves its upstream graph un-evaluated on exactly those cycles, so a
Materialize placed downstream of the binning is a no-op when it
matters most. Place it upstream of any stage that can withhold output.
Functions
- materialize_array(data, mode=MaterializeMode.SYNC)[source]#
Apply a
MaterializeModeto one array, and return it unchanged.A no-op for every backend but MLX: NumPy, CuPy and PyTorch arrays are already materialized by the time they are returned, so there is nothing to force. Returns
dataitself rather than a copy — the evaluation is a side effect on the array’s pending computation graph, not a new value.- Parameters:
data – The array to evaluate. Non-MLX arrays pass through untouched.
mode (MaterializeMode | str) – The
MaterializeModeto apply.
- Returns:
data, unchanged.- Raises:
ValueError – If
modeis not a validMaterializeMode.
Classes
- class Materialize(*args, settings=None, **kwargs)[source]#
Bases:
BaseTransformerUnit[MaterializeSettings,AxisArray,AxisArray,MaterializeTransformer]- Parameters:
settings (Settings | None)
- SETTINGS#
alias of
MaterializeSettings
- class MaterializeMode(*values)[source]#
-
How to force a lazy array backend to evaluate.
- SYNC = 'sync'#
Evaluate and block until the device finishes (
mx.eval).The strongest guarantee and the most expensive: the calling thread stalls for a full device round-trip, measured at ~0.1-0.15 ms per call on an M4 Pro regardless of array size. Use when the very next thing you do needs the values on the host, or when you are timing the work and want the cost attributed here rather than to whoever touches the array next.
- ASYNC = 'async'#
Schedule the evaluation and return immediately (
mx.async_eval).Detaches the computation graph just as
SYNCdoes — so it is equally effective at stopping a lazy graph from accumulating across calls — without the round-trip stall, leaving the CPU free to build the next message while the device works. Prefer it wherever the guarantee you want is “the graph does not grow” rather than “the values are on the host now”. Roughly 2xSYNCon a streaming chain; see the module docstring for where to put it.The trade is that it puts no bound on how far the producer may run ahead of the device, and the run-ahead is held as in-flight intermediates. Replaying 400 messages of 512x1024 float32 as fast as the host could build them: 2083 MiB in flight at 235 ms wall, against 33 MiB at 366 ms for
SYNC— 1.6x the throughput for 60x the transient memory. A live source that paces itself never gets near that; a file replay with no downstream backpressure and no other evaluation point can.
- OFF = 'off'#
Do nothing; leave the data lazy.
The fastest option and the correct one when something downstream — a conversion back to NumPy, an outlet, a serializer — is already going to force evaluation on every cycle. Leaves the node wired as an exact pass-through so the mode can be changed without editing the graph.
- class MaterializeSettings(mode: MaterializeMode = <MaterializeMode.SYNC: 'sync'>)[source]#
Bases:
Settings- Parameters:
mode (MaterializeMode)
- mode: MaterializeMode = 'sync'#
this node exists only to be an evaluation barrier, so wiring one and getting anything weaker would be surprising. Set
MaterializeMode.ASYNCto bound the graph without stalling, orMaterializeMode.OFFto disable the barrier while leaving it wired.- Type:
How to force evaluation. Defaults to
MaterializeMode.SYNC
- __init__(mode=MaterializeMode.SYNC)#
- Parameters:
mode (MaterializeMode)
- Return type:
None
- class MaterializeTransformer(*args, settings=None, **kwargs)[source]#
Bases:
BaseTransformer[MaterializeSettings,AxisArray,AxisArray]- Parameters:
settings (SettingsType)