feat: lower while and runtime calls#123
Merged
Merged
Conversation
De-risks native `:while` lowering (see EMLX.Native.Expr's `:while` limitation): a genuine mlx::core::Primitive (EMLXWhileSpike), pinned to the CPU stream like the existing linalg/runtime_call primitives, can call mlx::core::eval() in a loop from inside its own eval_cpu -- itself reached via eval_program's outer mlx::core::eval(outputs) on the worker's dedicated OS thread -- without deadlocking, even when the outer program's default device/stream is GPU. This validates the core mechanism a real EMLXWhile primitive needs: read back a loop condition scalar and re-invoke sub-computation natively in C++, once per eval_program call, instead of round-tripping to BEAM per iteration. Temporary: EMLXWhileSpike/:while_spike is not reachable from any real Elixir lowering path yet, only from the accompanying hand-built-Program test. Will be replaced by the real EMLXWhile primitive.
Add RefKind::Carry and an Instruction.subprograms field (decoded into a new SubProgram struct) so a future native :while instruction can carry self-contained cond/body instruction lists instead of splicing them into the parent program. Round-trip validated by a NIF-level test asserting the nested wire format decodes cleanly (reaching the "unknown op" check for the not-yet-implemented :while op registry entry). Checkpoint b of native while-loop lowering (see pv-feat/lower-while).
Add the EMLXWhile mlx::core::Primitive: interprets a :while instruction's cond/body SubPrograms directly in C++, iterating until cond is false, all inside a single eval_cpu call. Pinned to the CPU stream, matching the pattern validated by the earlier de-risking spike. Each iteration's carry is forced with mlx::core::eval() before feeding it back in, to avoid unbounded graph growth across a data-dependent iteration count. Refactors the flat instruction-list interpreter previously inlined in compile_program_impl's lambda into a shared interpret_instructions helper, reused for both a Program's top-level instructions and, via EMLXWhile, a :while instruction's nested cond/body sub-programs. Also extends compile_program_impl's op-name validation to recurse into :while sub-programs so unknown ops there are caught at compile time rather than eval time, and removes the now-superseded EMLXWhileSpike de-risking spike (its class, op_registry entry, and test) per its own "remove once the real EMLXWhile primitive lands" note. Validated by two new NIF-level tests: a single-carry-slot counting loop and a two-independent-carry-slot loop, both driving the primitive via a hand-built EMLX.Native.Program (no Elixir lowering path emits :while yet — that's checkpoint d). Checkpoint c of native while-loop lowering (see pv-feat/lower-while).
expr.ex: when a :while's trip count can't be statically unrolled,
lower it directly to a native :while wire instruction (EMLXWhile,
checkpoint c) instead of raising, provided native_while_eligible?/2
holds -- no :runtime_call, hook, or nested :while anywhere in the
condition/body (side effects and untested nested-eval reentrancy are
out of scope for now; such loops still fall back to the host-driven
path). Condition/body lower into isolated internal sub-programs
(lower_while_subprogram/4) with fresh carry refs standing in for the
while's own parameters; to_native/1 converts them into wire
EMLX.Native.SubPrograms with their own local {:result, i} numbering
and {:carry, i} refs, seeded only from the outer program's shared
capture/const tables (never its inputs or results) -- structurally
enforcing Nx.Defn.while's closure rules rather than just trusting them.
The shared per-instruction wire conversion is factored out of to_native/1
into to_native_instructions/4, reused recursively for sub-programs.
emlx.ex: split_point?/1 now only treats a while as a graph-split point
when it's native-ineligible; build_eval_fn's bare-while branch is
gated the same way, so an eligible bare tail while also takes the flat
native path instead of the host-driven one. Ineligible loops (including
nested whiles, whose *inner* while becomes eligible once the outer
loop's host-driven re-entry compiles its body standalone) keep using
Nx.Defn.Graph.split + run_while_loop exactly as before.
Also fixes a pre-existing inefficiency where a top-level while with a
static trip count always went through the split/host-driven path
instead of expr.ex's existing unroller, since detect_static_while_trip_count
now gets a chance to run before falling through to the native/fallback
decision.
Validated by new unit tests (native_while_eligible?/2 true/false cases,
a direct assertion that a dynamic while lowers to one :while wire
instruction) plus the existing while/hook/runtime_call/nested-while
end-to-end suites, none of which needed changes to keep passing.
Checkpoint d (final) of native while-loop lowering (see pv-feat/lower-while).
The wire format, Elixir lowering, and C++ interpreter were already written generically with no single-level-only assumption -- the only thing blocking a while-inside-while from lowering to a nested EMLXWhile primitive was the eligibility check's blanket exclusion. Validated safe (2-, 3-, and 20-level nesting; default GPU device; full existing suite) against checkpoint (a)'s reentrant-eval() spike, which generalizes here since native call-stack depth scales with lexical nesting depth, not iteration count. Added a defensive 64-level depth cap (advisor-recommended) purely to turn pathological/ generated deep nesting into a clean Elixir error instead of a native stack overflow.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR removes the need for graph splitting on runtime calls and while loops.
Benchmarks now stand as follows:
Qwen3-0.6B (dense bf16)
Qwen3-0.6B-MLX-4bit