MLX: bound the lazy graph so a method's peak memory is not its whole graph

MLX is lazy: Interpreter::dispatch only builds graph nodes, and nothing is
materialized until MLXBackend::execute calls async_eval on the method outputs.
For a long instruction chain that means every intermediate in the method is
live at the same instant.

Whisper-small encode is 495 instructions built before a single byte is
evaluated. Measured with mlx::core::get_{active,peak,cache}_memory on macOS:
peak 1105.6 MB against 94.8 MB of steady-state active memory. On iOS that peak
lands in the app's footprint and is what makes the model unusable there
(#22513).

Evaluate the live per-execution tensors once the intermediates produced since
the last barrier exceed a byte budget. Each barrier costs a GPU sync, so the
cost tracks the NUMBER of barriers; budgeting bytes rather than counting
instructions puts them only in the methods that actually allocate. Whisper-small
at 512 MB takes 12 barriers in encode and 0 in decode.

iPhone 16, whisper-small int8 through the full pipeline, medians of interleaved
rounds:

                     peak MB   peak-loaded   pipeline ms   encode   decode x10
  no barrier          1194.4       763.7          885.2     463.2        472.6
  byte budget 512MB    692.8       261.0          831.3     420.5        410.4

1.72x lower peak, 2.9x lower execute-phase footprint, and 1.06x faster than the
unbounded path.

macOS, medians of 4 interleaved rounds x 10 executions, round 1 discarded:

              peak MB (off -> 512)   speed vs off
  small encode   1105.6 -> 350.1        1.05x
  small decode    258.0 -> 258.0        0.97x
  tiny  encode    550.4 -> 343.2        1.07x
  tiny  decode     61.8 ->  61.8        0.93x
  SmolLM2 fwd    1196.5 -> 496.8        1.00x

Outputs bit-identical to the unbounded path on whisper tiny/base/small x
{encode, decode} and on two MLX LLMs. ET_MLX_EVAL_BUDGET_MB overrides the
budget; 0 restores the previous behaviour.

Author: msluszniak

backends/mlx/runtime/MLXInterpreter.h

@@ -19,6 +19,9 @@
 #include <mlx/mlx.h>
 #include <mlx/ops.h>
 
+#include <algorithm>
+#include <cstdlib>
+
 namespace executorch {
 namespace backends {
 namespace mlx {
@@ -1933,6 +1936,19 @@ class Interpreter {
     run_chain(prog, prog.main_chain_idx, st, stream);
   }
 
+  // Bytes of pending intermediates allowed to accumulate before a forced eval.
+  // 0 disables the barrier entirely (the pre-#22513 behaviour).
+  static size_t eval_budget_bytes() {
+    static const size_t bytes = [] {
+      const char* e = std::getenv("ET_MLX_EVAL_BUDGET_MB");
+      size_t mb = e == nullptr
+          ? 512u
+          : static_cast<size_t>(std::strtoul(e, nullptr, 10));
+      return mb * 1024u * 1024u;
+    }();
+    return bytes;
+  }
+
   void run_chain(
       const MLXProgram& prog,
       uint32_t chain_idx,
@@ -1945,6 +1961,27 @@ class Interpreter {
           std::to_string(prog.instruction_chains.size()) + ")");
     }
     const auto& chain = prog.instruction_chains[chain_idx];
+    // MLX is lazy: dispatch() only builds graph nodes, and nothing is
+    // materialized until MLXBackend::execute calls async_eval on the outputs.
+    // For a long chain that means every intermediate in the method is live at
+    // the same time. Whisper-small's 495-instruction encode peaks at 1105 MB of
+    // MLX allocation against 95 MB of steady-state active memory, which is what
+    // makes the model unusable on an iPhone (pytorch/executorch#22513).
+    //
+    // Bound it by evaluating once the intermediates produced since the last
+    // barrier exceed a byte budget. Each barrier costs a GPU sync, so the cost
+    // tracks the NUMBER of barriers, and the budget is on bytes rather than an
+    // instruction count so that only methods which actually allocate get any.
+    // Whisper-small at 512 MB takes 12 barriers in encode and 0 in decode,
+    // where an every-32-instruction rule took 15 and 22 -- and those 22 bought
+    // 50 MB on a method that peaks at 258 MB while costing 21% on an iPhone 16.
+    //
+    // Per instruction we add the largest tensor it touches, which tracks the
+    // size of what it just produced without needing to know which tid is the
+    // output. Evaluating early does not change results (verified
+    // bit-identical).
+    const size_t eval_budget = eval_budget_bytes();
+    size_t pending_bytes = 0;
     size_t idx = 0;
     for (const auto& instr : chain) {
       st.begin_op(idx, op_name(instr.op));
@@ -1957,6 +1994,32 @@ class Interpreter {
       }
       st.end_op();
       ++idx;
+
+      if (eval_budget != 0) {
+        size_t widest = 0;
+        for_each_tid(instr, [&](Tid id) {
+          if (id.idx >= st.num_constants && !st.is_mutable_buffer(id)) {
+            uint32_t slot = st.tensor_index(id);
+            if (slot < st.tensors.size() && st.tensors[slot].has_value()) {
+              widest = std::max(widest, st.tensors[slot]->nbytes());
+            }
+          }
+        });
+        pending_bytes += widest;
+        if (pending_bytes >= eval_budget) {
+          std::vector<::mlx::core::array> live;
+          live.reserve(st.tensors.size());
+          for (auto& t : st.tensors) {
+            if (t.has_value()) {
+              live.push_back(*t);
+            }
+          }
+          if (!live.empty()) {
+            ::mlx::core::eval(live);
+          }
+          pending_bytes = 0;
+        }
+      }
     }
   }