feat: control flow — :cond (where-chain) + honest deferrals (CM4) - #151
Merged
Conversation
- :cond — raw args [clauses, last]; lower to a select chain where(p1, b1, where(p2, b2, ... last)). All branches are evaluated (Nx branches are side-effect-free and shape-compatible) and the result matches the Evaluator's chosen branch exactly; only the not-taken branches' compute is wasted. The whole-tensor scalar predicate selects a branch wholesale. - :attach_token — pass through to the inner expr when there are no active hooks; raise clearly if hooks are present (a mid-graph Elixir callback isn't expressible in the single-NIF replay; program-split is deferred). Tests: compiler_control_flow — two-branch if, multi-clause cond, and nested conds, each bit-identical to the Evaluator. :while is next.
- reduce / window_reduce with an arbitrary BEAM reducer raise a clear error (no silent fallback); the fixed-identity aggregates (sum/product/ reduce_max/reduce_min) lower natively as their own ops. - while and its :elem tuple projection raise a clear "deferred" error — the single-NIF replay has no loop construct, so a data-dependent while needs static-trip unrolling or a worker-side synced loop (a focused follow-up). defn while is not used by the core transformer forwards; decode/generation loops run in Elixir. Tests assert both raise clearly under the native compiler.
- Split :elem out of the :while raise with an accurate message — :elem is a generic tuple projection (defn while AND multi-output ops), not while- specific, so the old "while loops" diagnostic could mislead. - Caveat the cond comment: a not-taken branch is still computed; on MLX an out-of-bounds gather there clamps (so the discarded value never changes the result), but a hard-faulting op on a not-taken path would diverge from the Evaluator's lazy single-branch eval.
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.
CM4 —
Nx.Defncontrol flow in the Expr→MLX compiler. Targetsfeat/expr-compiler.What this adds
:cond/if→ awhere-chainwhere(p1, b1, where(p2, b2, … last)), so a defn cond compiles single-NIF bit-identical to theEvaluator. First matching predicate wins (matches Nx semantics). All
branches are evaluated (Nx branches are side-effect-free + shape-
compatible) and the select picks the right one — only the not-taken
branches' compute is wasted (documented).
:attach_token→ pass-through to the inner expr when there are noactive hooks; raises clearly if hooks are present (a mid-graph Elixir
callback isn't expressible in the single-NIF replay).
reduce/window_reducewith an arbitrary BEAM reducer (the fixedaggregates lower natively);
whileand its:elemtuple projection(the replay has no loop construct — a focused follow-up; defn
whileisn't used by the core transformer forwards, whose loops run in Elixir).
Gate
compiler_control_flow_test— two-branchif, multi-clausecond,nested conds (each bit-identical to the Evaluator), plus assertions that
arbitrary
reduceandwhileraise clearly.mix precommitgreen:40 doctests, 79 properties, 615 tests, 0 failures; credo --strict
clean. Reviewed (focused); findings addressed.
Next
:while(static-unroll / worker-synced loop) as a follow-up; then CM5(no-fallback conformance) and CM6 (
mx::compilesecondary).