diff --git a/src/ezmsg/sigproc/asarray.py b/src/ezmsg/sigproc/asarray.py index 78a129d..a601688 100644 --- a/src/ezmsg/sigproc/asarray.py +++ b/src/ezmsg/sigproc/asarray.py @@ -80,6 +80,49 @@ def _detect_backend(data) -> str: raise TypeError(f"Unrecognized array type: {type(data)} (module={module})") +_MLX_CACHE_LIMIT_APPLIED: float | None = None +"""Value this process has already handed to ``mx.set_cache_limit``. + +Module-level because the limit is a property of the process, not of a node: two +``AsArray`` nodes converting to MLX in one process share one allocator. +""" + + +def _apply_mlx_cache_limit(limit_mb: float) -> None: + """Bound the MLX buffer cache for this process. Idempotent. + + MLX caches every freed buffer in a multimap keyed by *exact* byte size, and + only reuses one within ``min(2 * size, size + 2 * page_size)`` of the + request -- above ~32 KiB that is effectively an exact match. A graph whose + message length varies therefore mints a permanent new size class per length, + in a cache whose default limit is the whole machine (23 GiB on a 24 GiB + host). Measured on a 30 kHz feature chain that sees 40 distinct message + lengths over an hour: 6015 MiB of physical footprint unbounded, 966 MiB at + 512 MiB, with *higher* throughput at the limit because there is less memory + pressure. + + Eviction is LRU, which is what makes a limit the right tool rather than a + blunt one: the steady-state shape is re-touched every message and stays at + the head, while one-off shapes from a stall fall to the tail and are freed + first. Measured, the hot allocation is unaffected (0.97x) after 39 rare + shapes have been evicted past a 128 MiB limit. + """ + global _MLX_CACHE_LIMIT_APPLIED + if _MLX_CACHE_LIMIT_APPLIED == limit_mb: + return + import mlx.core as mx + + if _MLX_CACHE_LIMIT_APPLIED is not None: + ez.logger.warning( + f"MLX cache limit already set to {_MLX_CACHE_LIMIT_APPLIED} MiB in this process; " + f"overriding with {limit_mb} MiB. The limit is process-global, so the last AsArray " + "node to convert wins -- give every MLX-targeting AsArray in a process the same " + "mlx_cache_limit_mb." + ) + mx.set_cache_limit(int(limit_mb * 1024 * 1024)) + _MLX_CACHE_LIMIT_APPLIED = limit_mb + + class AsArraySettings(ez.Settings): backend: ArrayBackend = ArrayBackend.numpy """Target array backend.""" @@ -87,10 +130,43 @@ class AsArraySettings(ez.Settings): dtype: str | None = None """Target dtype as a string (e.g. "float32", "float64"). None keeps the original dtype.""" + mlx_cache_limit_mb: float | None = 512.0 + """Cap the MLX buffer cache (MiB) for the process that runs this node. + + Applied only when :attr:`backend` is MLX, once per process, on the first + message. ``None`` leaves MLX's default, which is the size of the machine. + + Sizing: one *distinct message shape* costs roughly **50x the message + payload** in cached buffers -- about 20 intermediates across a typical chain, + each keeping its own size class. Measured cache for a steady-state chain, + against ``samples x channels x 4`` bytes per message: 46x at 256 ch x 1200 + samples, 67x at 256 ch x 300. So:: + + limit_MiB ~= 50 * message_MiB * (distinct shapes to keep hot) + + A 256-channel, 300-sample float32 message is 0.29 MiB, so ~15 MiB per shape + and the 512 MiB default holds ~26 distinct shapes. Steady state needs only + one; the rest is headroom for the varying-length messages a stall produces. + Raise it if the graph legitimately cycles through many shapes; the floor is + one working set (~1x the ``50 *`` term), and 0 disables caching entirely at + a measured 40% throughput cost. + + This is a process-global MLX setting, so it is shared with anything else + using MLX in the same process. The default suits streaming graphs; large + offline batch work in the same process may want it raised or set to + ``None``.""" + class AsArrayTransformer(BaseTransformer[AsArraySettings, AxisArray, AxisArray]): def _process(self, message: AxisArray) -> AxisArray: target_backend = str(self.settings.backend) + if target_backend == "mlx" and self.settings.mlx_cache_limit_mb is not None: + # Here rather than in the Unit's initialize(): this transformer is + # also used bare (offline chains, benchmarks), and the limit has to + # be set in whichever process actually converts, which for a + # multi-process graph is not the one that built it -- set_cache_limit + # does not survive the spawn. + _apply_mlx_cache_limit(self.settings.mlx_cache_limit_mb) dtype_str = self.settings.dtype data = message.data diff --git a/src/ezmsg/sigproc/util/mlx_metal_common.py b/src/ezmsg/sigproc/util/mlx_metal_common.py index a24af51..b4e2292 100644 --- a/src/ezmsg/sigproc/util/mlx_metal_common.py +++ b/src/ezmsg/sigproc/util/mlx_metal_common.py @@ -72,11 +72,29 @@ def chunked_scan(x_flat, n_samples, chunk_sizes, state, launch_fn): ``launch_fn``. Kernels must emit their state at ``valid_length - 1`` rather than after the padding. + Chunk sizes bound the set of Metal kernel *specializations*; they do not by + themselves bound the set of *buffer sizes*. MLX caches freed buffers in a + multimap keyed by exact byte size and only reuses one within + ``min(2 * size, size + 2 * page_size)`` of the request -- effectively an + exact match above ~32 KiB -- so every distinct intermediate length becomes a + permanent new size class in a cache whose default limit is the whole + machine. Concatenating the *padded* chunks and trimming once therefore + allocates on the multiple-of-chunk_size grid rather than at ``n_samples``, + which measured 31% less cached memory over 40 distinct input lengths. + + Trimming once at the end is only correct because padding can occur on the + final chunk alone: every earlier iteration has ``remaining > chunk_size``, + so ``valid == chunk_size`` and the chunk is emitted whole. Were an interior + chunk padded, its padding would sit between two runs of valid samples and + the single trim would return garbage, so the loop asserts that invariant + rather than trusting whoever next edits the size-selection rule. + Returns ``(y_combined, final_state)``. """ y_chunks = [] start = 0 max_chunk_size = chunk_sizes[-1] + padded_len = 0 while start < n_samples: remaining = n_samples - start chunk_size = next((size for size in chunk_sizes if size >= remaining), max_chunk_size) @@ -84,10 +102,18 @@ def chunked_scan(x_flat, n_samples, chunk_sizes, state, launch_fn): end = start + valid x_chunk = x_flat[:, start:end] if valid < chunk_size: + if end != n_samples: + raise AssertionError( + f"chunked_scan padded an interior chunk (valid={valid}, chunk_size={chunk_size}, " + f"end={end}, n_samples={n_samples}); the single trim below would return padding as data." + ) x_chunk = mx.pad(x_chunk, [(0, 0), (0, chunk_size - valid)]) valid_length = mx.array([valid], dtype=mx.uint32) y_chunk, state = launch_fn(x_chunk, state, chunk_size, valid_length) - y_chunks.append(y_chunk[:, :valid]) + y_chunks.append(y_chunk) + padded_len += chunk_size start = end - y_combined = y_chunks[0] if len(y_chunks) == 1 else mx.concatenate(y_chunks, axis=-1) + y_padded = y_chunks[0] if len(y_chunks) == 1 else mx.concatenate(y_chunks, axis=-1) + # Exact length is part of the contract: callers reshape to n_samples. + y_combined = y_padded if padded_len == n_samples else y_padded[:, :n_samples] return y_combined, state diff --git a/tests/unit/test_asarray.py b/tests/unit/test_asarray.py index f31b258..cb1965c 100644 --- a/tests/unit/test_asarray.py +++ b/tests/unit/test_asarray.py @@ -13,6 +13,7 @@ _get_backend_module, ) from tests.helpers.empty_time import check_empty_result, make_empty_msg, make_msg +from tests.helpers.util import requires_mlx # -- Helpers ------------------------------------------------------------------ @@ -187,3 +188,93 @@ def test_empty_time_cross_backend(backend): # The time dimension should still be 0. time_idx = result.dims.index("time") assert result.data.shape[time_idx] == 0 + + +# -- MLX buffer cache limit ---------------------------------------------------- + + +def test_mlx_cache_limit_default_is_set(): + """A default, not None: the growth it prevents is invisible in RSS. + + MLX buffers are IOKit allocations, so a graph whose message length varies + can put gigabytes into the allocator's per-size cache while RSS stays flat. + Users do not go looking for a knob they have no symptom for. + """ + assert AsArraySettings().mlx_cache_limit_mb == 512.0 + + +def test_mlx_cache_limit_not_applied_for_numpy_target(monkeypatch): + """Converting TO numpy must not touch a process-global MLX setting.""" + import ezmsg.sigproc.asarray as asarray_module + + calls = [] + monkeypatch.setattr(asarray_module, "_apply_mlx_cache_limit", lambda mb: calls.append(mb)) + proc = AsArrayTransformer(AsArraySettings(backend=ArrayBackend.numpy, mlx_cache_limit_mb=128.0)) + proc(make_msg()) + assert calls == [] + + +@requires_mlx +def test_mlx_cache_limit_applied_once_for_mlx_target(monkeypatch): + import ezmsg.sigproc.asarray as asarray_module + + calls = [] + monkeypatch.setattr(asarray_module, "_apply_mlx_cache_limit", lambda mb: calls.append(mb)) + proc = AsArrayTransformer(AsArraySettings(backend=ArrayBackend.mlx, mlx_cache_limit_mb=128.0)) + for _ in range(3): + proc(make_msg()) + # Called per message, but the applier itself is what dedupes -- see below. + assert calls == [128.0, 128.0, 128.0] + + +@requires_mlx +def test_apply_mlx_cache_limit_is_idempotent_and_warns_on_conflict(monkeypatch): + """The limit is process-global, so a second, different value is a conflict. + + Silently letting the last node win would make the effective limit depend on + which unit happened to convert first. + """ + import mlx.core as mx + + import ezmsg.sigproc.asarray as asarray_module + + sets = [] + monkeypatch.setattr(mx, "set_cache_limit", lambda n: sets.append(n)) + monkeypatch.setattr(asarray_module, "_MLX_CACHE_LIMIT_APPLIED", None) + warnings = [] + monkeypatch.setattr(asarray_module.ez.logger, "warning", lambda msg, *a: warnings.append(msg)) + + asarray_module._apply_mlx_cache_limit(128.0) + asarray_module._apply_mlx_cache_limit(128.0) # same value: no-op + assert sets == [128 * 1024 * 1024] + assert warnings == [] + + asarray_module._apply_mlx_cache_limit(256.0) # different value: warn, override + assert sets == [128 * 1024 * 1024, 256 * 1024 * 1024] + assert len(warnings) == 1 and "process-global" in warnings[0] + + +@requires_mlx +def test_mlx_cache_limit_actually_bounds_the_cache(): + """End to end: the setting reaches the allocator and caps it.""" + import mlx.core as mx + + import ezmsg.sigproc.asarray as asarray_module + + previous = mx.set_cache_limit(2**40) + applied = asarray_module._MLX_CACHE_LIMIT_APPLIED + try: + asarray_module._MLX_CACHE_LIMIT_APPLIED = None + proc = AsArrayTransformer(AsArraySettings(backend=ArrayBackend.mlx, mlx_cache_limit_mb=32.0)) + proc(make_msg()) + mx.clear_cache() + # Churn many distinct sizes; without a limit this cache grows unbounded. + for n in range(1, 60): + a = mx.zeros((256, 300 * n)) + mx.eval(a) + del a + assert mx.get_cache_memory() <= 32 * 1024 * 1024 + finally: + mx.clear_cache() + mx.set_cache_limit(previous) + asarray_module._MLX_CACHE_LIMIT_APPLIED = applied diff --git a/tests/unit/test_mlx_metal_common.py b/tests/unit/test_mlx_metal_common.py new file mode 100644 index 0000000..518a0a8 --- /dev/null +++ b/tests/unit/test_mlx_metal_common.py @@ -0,0 +1,115 @@ +import numpy as np + +from tests.helpers.util import requires_mlx + + +@requires_mlx +def test_chunked_scan_trims_to_exact_length(): + """The padded-concat optimization must not change the output contract. + + ``chunked_scan`` concatenates the kernels' *padded* chunks so the allocation + lands on a multiple-of-chunk_size grid rather than at ``n_samples`` (fewer + MLX buffer size classes). The single trim at the end is what keeps the + result exact, and callers reshape to ``n_samples`` immediately afterwards -- + so an off-by-one here surfaces as a reshape error, not as wrong data. + """ + import mlx.core as mx + + from ezmsg.sigproc.util.mlx_metal_common import chunked_scan + + def launch(x_chunk, state, cs, valid_length): + # A stand-in kernel with the real ones' contract: output is the FULL + # chunk width, state carries forward. + assert x_chunk.shape[-1] == cs + return x_chunk * 2.0, state + 1 + + n_channels = 3 + for chunk_sizes in ((32,), (32, 128), (16, 64, 256)): + for n_samples in (1, 15, 16, 17, 31, 32, 33, 100, 128, 129, 257, 1000): + x = mx.array(np.arange(n_channels * n_samples, dtype=np.float32).reshape(n_channels, n_samples)) + y, state = chunked_scan(x, n_samples, chunk_sizes, mx.array(0), launch) + assert y.shape == (n_channels, n_samples), f"chunk_sizes={chunk_sizes} n_samples={n_samples} gave {y.shape}" + # Padding must never leak into the result. + np.testing.assert_allclose(np.asarray(y), np.asarray(x) * 2.0) + + +@requires_mlx +def test_chunked_scan_pads_only_the_final_chunk(): + """Trimming once is sound only while padding is confined to the last chunk. + + Padding an interior chunk would put padding *between* two runs of valid + samples and the single trim would return it as data. The invariant holds + structurally -- padding needs ``remaining < chunk_size``, which makes that + chunk the tail -- so this pins the property across size sets and lengths + rather than trying to construct a violation (there isn't one to construct; + the assertion in ``chunked_scan`` guards future edits to the selection rule). + """ + import mlx.core as mx + + from ezmsg.sigproc.util.mlx_metal_common import chunked_scan + + for chunk_sizes in ((32,), (32, 128), (16, 64, 256), (128, 32)): + for n_samples in (1, 15, 31, 32, 33, 100, 128, 129, 257, 1000): + seen = [] + + def launch(x_chunk, state, cs, valid_length, _seen=seen): + _seen.append((int(np.asarray(valid_length)[0]), cs)) + return x_chunk, state + + chunked_scan(mx.zeros((2, n_samples)), n_samples, chunk_sizes, mx.array(0), launch) + padded = [i for i, (valid, cs) in enumerate(seen) if valid < cs] + assert padded in ([], [len(seen) - 1]), ( + f"chunk_sizes={chunk_sizes} n_samples={n_samples}: padded chunks at {padded} of {len(seen)}" + ) + + +@requires_mlx +def test_chunked_scan_allocates_fewer_size_classes_than_per_chunk_trim(): + """The point of the change, measured against the implementation it replaced. + + Compares like with like: identical inputs, held live so their own size + classes are not what is being counted, so the delta is the scan's + intermediates alone. A MiB figure would be machine-dependent; the ratio is + the property worth pinning. + """ + import mlx.core as mx + + from ezmsg.sigproc.util.mlx_metal_common import chunked_scan + + chunk_sizes = (32, 1024) + + def launch(x_chunk, state, cs, valid_length): + return x_chunk * 2.0, state + + def per_chunk_trim(x_flat, n_samples, sizes, state, launch_fn): + """The pre-change behavior: trim every chunk, concatenate at n_samples.""" + y_chunks, start = [], 0 + while start < n_samples: + remaining = n_samples - start + cs = next((s for s in sizes if s >= remaining), sizes[-1]) + valid = min(remaining, cs) + x_chunk = x_flat[:, start : start + valid] + if valid < cs: + x_chunk = mx.pad(x_chunk, [(0, 0), (0, cs - valid)]) + y, state = launch_fn(x_chunk, state, cs, mx.array([valid], dtype=mx.uint32)) + y_chunks.append(y[:, :valid]) + start += valid + y = y_chunks[0] if len(y_chunks) == 1 else mx.concatenate(y_chunks, axis=-1) + return y, state + + lengths = [300 * k for k in range(1, 21)] * 3 + inputs = [mx.zeros((256, n)) for n in lengths] # held live throughout + mx.eval(*inputs) + + def cache_for(scan): + mx.clear_cache() + for x, n in zip(inputs, lengths): + y, _ = scan(x, n, chunk_sizes, mx.array(0), launch) + mx.eval(y) + used = mx.get_cache_memory() + mx.clear_cache() + return used + + before = cache_for(per_chunk_trim) + after = cache_for(chunked_scan) + assert after < before, f"padded concat cached more, not less: {after} vs {before} bytes"