feat: Expr→MLX single-NIF compiler — lowerer + core primitive coverage (CM1) - #147
Merged
Conversation
Grow the Expr-compiler IR wire format with an `iattrs` field (integer attributes: dtype codes, shapes, axes) so ops beyond elementwise can be replayed, and add the elementwise / cast / shape opcode set: - C++: opcodes.hpp gains ~40 opcodes (binary arith/compare/logical/ bitwise, unary, astype/reshape/transpose/squeeze/broadcast_to) with a dispatch_op switch calling the same mlx::core::* entry points as the eager NIFs; dtype codes decoded via to_mlx_dtype_code. compile_program / eval_program / describe_program thread iattrs through. - Elixir: Emily.IR gains the full opcode table, dtype_code/1, and per-instruction :iattrs; Emily.Program ships them. Tests: hand-built IR replay-vs-eager for unary+cast+binary chains, reshape/transpose, broadcast_to, plus an iattrs describe round-trip. The wire-format change updates the CM0 call sites. The Expr->IR lowerer and Emily.Compiler integration follow next.
Lower a traced Nx.Defn.Expr to the flat IR and run it through the program-replay engine in one NIF call, behind Emily.Compiler's new `native: true` option (the Evaluator path stays the default). - Emily.IR.lower/1 walks the Expr DAG (memoized by node id), porting Emily.Backend's per-op logic incl. dtype coercion: arithmetic casts operands to out.type; compares cast to merge(a,b) then coerce bool->u8; shape/broadcast mirror the backend's reshape + broadcast_to. Parameters become input slots; constants and tensor literals are materialized as captured refs. Unsupported ops raise (no silent fallback). - Emily.Compiler gains compile_native/2: trace -> lower -> compile a Program once (cached in the closure) -> replay per call, realizing param thunks to refs and reassembling the output container. Tests: compiler_equivalence_test runs each function via the native path vs the Evaluator and asserts bit-identical across unary, binary (incl. broadcasting + mixed-dtype coercion), compare, cast/shape, constants, tensor captures, DAG sharing, tuple outputs, identity, and reuse. Coverage so far is elementwise + cast + shape; dot, reductions, gather/scatter and the remaining primitives follow in later CM1 commits toward the Axon MLP / DistilBERT gate.
Add matmul/tensordot and reduction (sum/product/max/min/all/any) opcodes and lower the corresponding Nx ops, porting Emily.Backend's logic: - dot: non-batched -> tensordot over the contraction axes; batched -> permute/flatten to 3-D, MLX matmul, reshape back (mirrors Backend.batched_matmul/7). Non-float batched dot raises (MLX matmul is float-only; the compiler does not fall back). - reductions: axes default to all (Nx.axes/1), keepdims threaded via iattrs; result coerced to the node's out.type like Backend.wrap. Tests: compiler_equivalence cases for 2D / vector / tensordot / batched dot, sum/product/max/min with keep_axes, a softmax-style exp/sum/divide composite, and a full two-layer MLP forward — all bit-identical to the Evaluator. The matmul-dominated path now runs single-NIF.
- select -> cast pred to {:pred, 1} then where (mirrors Backend.select/4).
- slice -> mx::slice with static integer starts; stops = starts+lengths.
Dynamic (tensor) starts raise (deferred to CM3's offset-as-input).
- iota -> materialized as a captured constant (shape/axis/type static).
Equivalence tests: select/where (+ relu via select), static slice (incl.
the [[range]] sugar), and iota — all bit-identical to the Evaluator.
- Stop tracking .claude/scheduled_tasks.lock (a machine/run-specific runtime lock) and add it to .gitignore. - Fix the stale Emily.Program.describe/1 @spec/@doc — it now returns the 7-tuple (with iattrs) the NIF produces. - Coerce the unary-op output to the node's out.type, matching Emily.Backend.wrap/3 (every eager op is wrapped through coerce). It's a no-op astype on aligned dtypes but keeps the node dtype exact and consistent with the reduction/dot/select/slice clauses. - Dedupe the constant/tensor/iota materialization into materialize_const/2 and materialize_capture/2 so the from_binary + list-prepend + index-bump can't desync across clauses.
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.
Second milestone (CM1) of the
Nx.Defn.Expr→ MLX single-NIF compiler,on top of CM0's program-replay engine. Targets
feat/expr-compiler.What this adds
A real
Nx.Defn.Expr→ flat IR lowerer and theEmily.Compilernative: truepath: a traced defn is lowered once, compiled into oneProgram, and replayed in a single NIF call per invocation — instead ofone BEAM↔worker round-trip per op (the Evaluator walk). The Evaluator
path stays the default; the compiled path is opt-in.
Emily.IR.lower/1walks the Expr DAG (memoized by node id),porting
Emily.Backend's per-op logic incl. dtype coercion. Parametersbecome input slots; constants / tensor literals / iota are materialized
as captured refs. Unsupported ops raise — no silent fallback.
iattrs: dtypecodes, shapes, axes);
compile_program/eval_program/describe_programthread them.
bitwise), cast, shape (reshape/transpose/squeeze/broadcast), dot
(matmul/tensordot + batched), reductions (sum/product/max/min/all/
any), select/where, static slice, iota.
Correctness gate
compiler_equivalence_testruns each function two ways — the nativesingle-NIF path vs
Nx.Defn.EvaluatorthroughEmily.Backend— andasserts bit-identical output. 27 cases cover unary, binary (incl.
broadcasting + mixed-dtype coercion), compare (bool→u8), cast/shape,
dot (2D/vector/tensordot/batched), reductions + softmax composite,
select/relu, slice, iota, a two-layer MLP forward, DAG sharing,
tuple outputs, and closure reuse. Plus per-opcode replay-vs-eager and IR
round-trip tests.
mix precommitgreen: 40 doctests, 79 properties, 596 tests, 0failures; credo --strict clean.
Review
Ran a multi-angle
/code-review; addressed findings in093e3ef(untrack a stray
.claudelock file + gitignore it; fix the staledescribe/1spec; coerce unary output to out.type for fullwrap-parity; dedupe the materialization blocks).Scope
CM1 covers the core primitive set + MLP single-NIF (no fallback). Ops a
full DistilBERT forward additionally needs (gather/embeddings, layernorm
neighbourhood, etc.) fold into later coverage work; fused kernels +
quantized_matmul are CM2. No public API or decoder numerics change (the
native path is opt-in).