[ExecuTorch][MLX] lowering memory/export improvements

backends/mlx/_memprofile.py

@@ -0,0 +1,188 @@
+#
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+"""Phase-attributed peak-memory profiling for MLX export.
+
+Enabled with ``ET_MLX_MEM_PROFILE=1``; inert otherwise::
+
+    ET_MLX_MEM_PROFILE=1 python -m ...export_dflash --target-gguf ...
+
+    [mem] build: 42.84 GB -> 41.30 GB, -1.54 GB net, high-water 48.26 GB, \
+RAISED PEAK by 5.42 GB, 24.4s
+
+Reports macOS *physical footprint* -- what ``/usr/bin/time -l`` calls "peak
+memory footprint" and Activity Monitor shows as "Memory". RSS is not a usable
+proxy: a 16 GB model export peaked at 65 GB footprint versus 37 GB RSS.
+
+Only phases tagged ``RAISED PEAK`` can lower the number ``/usr/bin/time -l``
+reports; everything else is churn under the existing high-water mark.
+
+The watermark comes from the kernel rather than a sampling thread, which would
+miss short-lived spikes -- a 2 GB allocation freed within 100 ms went unseen at
+a 50 ms sampling interval.
+"""
+
+from __future__ import annotations
+
+import ctypes
+import ctypes.util
+import os
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+from executorch.backends.mlx._logging import logger
+
+_ENV_VAR = "ET_MLX_MEM_PROFILE"
+
+# proc_pid_rusage(pid, flavor, rusage_info_t *buffer), from <libproc.h>.
+# Offsets into struct rusage_info_v4 (<sys/resource.h>): a 16-byte ri_uuid
+# followed by uint64 counters. Sanity-checked at runtime by _self_check().
+_RUSAGE_INFO_V4 = 4
+_OFF_PHYS_FOOTPRINT = 72
+_OFF_PROC_START_ABSTIME = 80
+_OFF_PROC_EXIT_ABSTIME = 88
+_OFF_LIFETIME_MAX_PHYS_FOOTPRINT = 240
+_RUSAGE_BUF_BYTES = 512  # generous; larger than any rusage_info_v4
+
+# Bounds for the self-check: any live process is above the floor, and no real
+# footprint approaches the ceiling.
+_MIN_PLAUSIBLE = 1 << 20  # 1 MiB
+_MAX_PLAUSIBLE = 1 << 50  # 1 PiB
+
+_libc = None
+_usable: Optional[bool] = None
+_depth = 0
+
+
+def _read_rusage() -> Optional[bytes]:
+    global _libc
+    try:
+        if _libc is None:
+            _libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)
+        buf = (ctypes.c_uint8 * _RUSAGE_BUF_BYTES)()
+        if _libc.proc_pid_rusage(os.getpid(), _RUSAGE_INFO_V4, ctypes.byref(buf)) != 0:
+            return None
+        return bytes(buf)
+    except (AttributeError, OSError, ValueError):
+        return None
+
+
+def _field(raw: Optional[bytes], offset: int) -> Optional[int]:
+    if raw is None:
+        return None
+    return int.from_bytes(raw[offset : offset + 8], "little")
+
+
+def _self_check() -> bool:
+    """Smoke-test the hardcoded rusage offsets before trusting what they return.
+
+    The struct layout is pinned by the flavor argument (Apple adds rusage_info_v5
+    rather than reordering v4), so this is a guard against being wrong about the
+    layout, not against it changing underneath us.
+
+    ri_proc_exit_abstime is the anchor: it is necessarily zero for the process
+    asking, so a non-zero read means the offsets are not landing where we think.
+    The bounds checks on the two footprint fields are weaker -- a misread can
+    still satisfy them by coincidence, since neighbouring counters hold numbers
+    of a similar magnitude -- so they catch gross errors only.
+    """
+    raw = _read_rusage()
+    current = _field(raw, _OFF_PHYS_FOOTPRINT)
+    peak = _field(raw, _OFF_LIFETIME_MAX_PHYS_FOOTPRINT)
+    started = _field(raw, _OFF_PROC_START_ABSTIME)
+    exited = _field(raw, _OFF_PROC_EXIT_ABSTIME)
+
+    if current is None:
+        reason = "proc_pid_rusage unavailable (non-macOS?)"
+    elif exited != 0 or not started:
+        reason = (
+            f"struct layout unrecognized (proc_start={started}, proc_exit={exited}; "
+            "expected a non-zero start and a zero exit for a live process)"
+        )
+    elif not _MIN_PLAUSIBLE <= current <= peak <= _MAX_PLAUSIBLE:
+        reason = f"implausible readings (current={current}, lifetime max={peak})"
+    else:
+        return True
+
+    logger.warning(f"[mem] memory profiling disabled: {reason}")
+    return False
+
+
+def enabled() -> bool:
+    """Whether profiling is switched on and the platform counters are trustworthy."""
+    global _usable
+    if os.environ.get(_ENV_VAR, "0") == "0":
+        return False
+    if _usable is None:
+        _usable = _self_check()
+    return _usable
+
+
+def phys_footprint() -> Optional[int]:
+    """Current physical footprint in bytes, or None if unavailable."""
+    return _field(_read_rusage(), _OFF_PHYS_FOOTPRINT)
+
+
+def peak_footprint() -> Optional[int]:
+    """Process lifetime maximum physical footprint in bytes.
+
+    Matches the "peak memory footprint" line from ``/usr/bin/time -l``.
+    """
+    return _field(_read_rusage(), _OFF_LIFETIME_MAX_PHYS_FOOTPRINT)
+
+
+def _gb(n: Optional[int]) -> str:
+    return "?" if n is None else f"{n / (1 << 30):.2f} GB"
+
+
+def _delta(before: Optional[int], after: Optional[int]) -> Optional[int]:
+    return None if (before is None or after is None) else after - before
+
+
+@contextmanager
+def mem_phase(name: str):
+    """Log footprint across `name`, attributing any new high-water mark to it.
+
+    Nested phases are indented. Does nothing unless ET_MLX_MEM_PROFILE is set.
+    """
+    global _depth
+
+    if not enabled():
+        yield
+        return
+
+    indent = "  " * _depth
+    start, start_peak = phys_footprint(), peak_footprint()
+    started_at = time.perf_counter()
+    _depth += 1
+    try:
+        yield
+    finally:
+        _depth -= 1
+        end, end_peak = phys_footprint(), peak_footprint()
+        net = _delta(start, end)
+        raised = _delta(start_peak, end_peak)
+
+        sign = "+" if net is not None and net >= 0 else ""
+        message = (
+            f"[mem]{indent} {name}: {_gb(start)} -> {_gb(end)}, "
+            f"{sign}{_gb(net)} net, high-water {_gb(end_peak)}"
+        )
+        if raised:
+            message += f", RAISED PEAK by {_gb(raised)}"
+        logger.info(f"{message}, {time.perf_counter() - started_at:.1f}s")
+
+
+def log_footprint(label: str) -> None:
+    """Log a one-off footprint reading."""
+    if enabled():
+        logger.info(
+            f"[mem] {label}: {_gb(phys_footprint())} "
+            f"(high-water {_gb(peak_footprint())})"
+        )

backends/mlx/builder/op_helpers.py

@@ -534,6 +534,28 @@ def emit_quantized_gather(
     )
 
 
+def mlx_qparams_supported(
+    in_features: int, num_groups: int, group_size: int, bits: int
+) -> bool:
+    """Whether to_mlx_qparams + regroup_affine_scales can consume this layout.
+
+    Mirrors their asserts using shape metadata only, so a handler can answer
+    "would I lower this?" without reading (and repacking) the weight:
+
+    * to_mlx_qparams packs a row into whole uint32 words, so
+      ``in_features * bits`` must be a multiple of 32.
+    * regroup_affine_scales repeat-interleaves, which only ever splits a group
+      finer, so the weight's own group must be a whole multiple of the
+      MLX-legal ``group_size``.
+    """
+    if in_features <= 0 or num_groups <= 0 or in_features % num_groups != 0:
+        return False
+    if (in_features * bits) % 32 != 0:
+        return False
+    weight_group_size = in_features // num_groups
+    return weight_group_size >= group_size and weight_group_size % group_size == 0
+
+
 def to_mlx_qparams(
     qdata: torch.Tensor,
     scale: torch.Tensor,

backends/mlx/builder/op_registry.py

@@ -8,7 +8,17 @@
 
 from __future__ import annotations
 
-from typing import Callable, Dict, List, Optional, Tuple, Type, TYPE_CHECKING, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    Type,
+    TYPE_CHECKING,
+    Union,
+)
 
 from executorch.backends.mlx._logging import logger
 from torch.fx.node import Node
@@ -23,6 +33,10 @@
     ["MLXProgramBuilder", Node], Optional[Union["Slot", Tuple["Slot", ...]]]
 ]
 
+# Support-check type: takes (builder, node) and returns whether the handler for
+# that node will lower it. See PatternHandler.supported for the contract.
+SupportCheck = Callable[["MLXProgramBuilder", Node], bool]
+
 
 class PatternHandler:
     def __init__(self, head: Node, body: List[Node]) -> None:
@@ -40,6 +54,32 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional[PatternHandle
     def __call__(self, P: MLXProgramBuilder, n: Node) -> None:
         raise NotImplementedError
 
+    def supported(self, P: MLXProgramBuilder, n: Node) -> bool:
+        """Optional: answer "would __call__ lower this node?" without running it.
+
+        Support is normally decided by running the handler and seeing whether it
+        throws, but handlers that repack a quantized weight copy the whole tensor
+        to answer that question -- and the partitioner asks it several times per
+        export. Overriding this lets such a handler answer from node metadata
+        instead; it is then consulted in place of running the handler during
+        support checks (never during build(), which has to emit).
+
+        Implementations must decide from ``n.meta`` alone and must never read
+        constant data, which is the cost being avoided. They must also agree with
+        __call__: saying yes to a node the handler then rejects makes
+        ops_to_not_decompose preserve an op that afterwards neither decomposes
+        nor lowers. Disagreement is caught by the op tests rather than a
+        dedicated check -- a false positive fails the export outright when
+        build() runs the handler for real, and a false negative shows up as a
+        missing delegate segment.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def has_support_check(cls) -> bool:
+        """Whether this class overrides supported()."""
+        return cls.supported is not PatternHandler.supported
+
     def set_handlers(self, P: MLXProgramBuilder):
         if P.node_info[self.head].handler is not None:
             raise AssertionError(
@@ -67,11 +107,13 @@ class MLXOpRegistry:
 
     def __init__(self):
         self._handlers: Dict[Union[str, Callable], Handler] = {}
+        self._support_checks: Dict[Union[str, Callable], SupportCheck] = {}
         self._patterns: Dict[str, Type[PatternHandler]] = {}
 
     def reset(self) -> None:
         """Reset the registry to empty state. Useful for testing."""
         self._handlers.clear()
+        self._support_checks.clear()
         self._patterns.clear()
 
     def register(self, target: Union[str, Callable, list, tuple]):
@@ -89,16 +131,44 @@ def deco(fn: Handler):
 
     def get_handler(self, node: Node) -> Optional[Handler]:
         """Get the handler for a node, or None if not registered."""
+        return self._lookup(self._handlers, node)
+
+    def register_support_check(self, target: Union[str, Callable, list, tuple]):
+        """Decorator registering a cheap support predicate for an op handler.
+
+        The predicate takes (builder, node) and returns whether the handler will
+        lower the node. It is consulted instead of running the handler during
+        support checks, so it must decide from node metadata alone -- reading
+        constant data is exactly the cost being avoided. See
+        PatternHandler.supported for the full contract.
+        """
+
+        def deco(fn: SupportCheck):
+            targets = target if isinstance(target, (list, tuple)) else [target]
+            for t in targets:
+                if t in self._support_checks:
+                    raise ValueError(f"Support check for {t} already registered")
+                self._support_checks[t] = fn
+            return fn
+
+        return deco
+
+    def get_support_check(self, node: Node) -> Optional[SupportCheck]:
+        """Get the support predicate for a node, or None if it has none."""
+        return self._lookup(self._support_checks, node)
+
+    @staticmethod
+    def _lookup(table: Dict[Union[str, Callable], Any], node: Node) -> Optional[Any]:
         t = node.target
-        if t in self._handlers:
-            return self._handlers[t]
+        if t in table:
+            return table[t]
         # Handle EdgeOpOverload by extracting the underlying ATen op
-        if hasattr(t, "_op") and t._op in self._handlers:
-            return self._handlers[t._op]
+        if hasattr(t, "_op") and t._op in table:
+            return table[t._op]
         # Check for string-based targets (e.g., higher_order ops)
         target_str = str(t)
-        if target_str in self._handlers:
-            return self._handlers[target_str]
+        if target_str in table:
+            return table[target_str]
         return None
 
     def registered_ops(self) -> set:
@@ -118,6 +188,8 @@ def unregister(self, target: Union[str, Callable, list, tuple]) -> None:
         for t in targets:
             if t in self._handlers:
                 del self._handlers[t]
+            if t in self._support_checks:
+                del self._support_checks[t]
 
     def register_pattern(self, name: str):
         """Decorator to register a pattern handler class."""