Native window/pooling ops — CNN forward + backward lowering - #179
Merged
Conversation
Add a `mode_test/2` macro to `Emily.ConformanceHelper` that emits each conformance test in three lanes from a single body, varying only the bound `predict_opts`: * eval — `[]`, the evaluator path (unchanged behaviour) * native — `native: true, native_fallback: :raise` * fusion — additionally `native_compiled: true` (mx::compile) The native and fusion lanes are tagged `:native` and `:native_compiled` on top of the module's `:conformance`, so `--only conformance` runs all three while `--only native` / `--only native_compiled` select one lane. Both compiler lanes resolve the same HuggingFace repos as the evaluator lane, so the checkpoint download is amortised across them. Convert the six tiny-random suites (DistilBERT, ViT, Whisper, ModernBert, NomicBert, SmolLm3). `Axon.predict`-based tests thread `predict_opts` into the call; `Axon.build`-based smoke tests build `init_fn` on the evaluator and only `predict_fn` under `predict_opts`, so the native gate covers the forward pass rather than random param init. The serving and fast-kernels tests stay eval-only. `native_fallback: :raise` makes the native lanes a no-fallback op-coverage gate: every architecture lowers fully, and the fusion lane holds at the existing 1e-4 reference tolerance (mx::compile only reassociates f32, so its drift is far below the approximate mx::fast::* kernels).
Extend `mode_test/2` to `mode_test/3` with a `lane_tags:` option and run the `*_full` forward-pass conformance suites (ViT-Base, Whisper-tiny) through the native and fusion compilers alongside the evaluator. The full suites pass `lane_tags: false`, so their native and fusion lanes are emitted without the cross-cutting `:native` / `:native_compiled` tags and stay gated behind the suite's own `:vit_full` / `:whisper_full` moduletag. `--only vit_full` now runs all three lanes on the full checkpoint, while `--only native` stays tiny-random only and never pulls a full-size download. Both full forwards lower fully under `native_fallback: :raise` and hold at the pinned 1e-4 reference tolerance in the fusion lane too — `mx::compile` only reassociates f32, so its drift stays well below the approximate `mx::fast::*` kernels (which need 1e-3 ViT / 1e-2 Whisper in the `:fast_kernels_full` variants). The serving and fast-kernels tests remain eval-only.
`Nx.take_along_axis` (the `Nx.Block.TakeAlongAxis` block) now lowers under `compiler: Emily.Compiler, native: true` instead of raising. The MLX op and eager NIF already existed (`Emily.Native.take_along_axis` → `mlx::core::take_along_axis`, used by `Emily.Backend.native_take_along_axis/4`); only the IR path was missing. - ir.ex: add opcode `take_along_axis` (81) and a `lower_block/5` clause mirroring the eager twin — cast indices to s32, emit, coerce to output type — alongside the existing `Nx.Block.Take` lowering. - opcodes.hpp: add `TakeAlongAxis = 81`, bump `kOpcodeCount`, and a dispatch case calling `mx::take_along_axis(in[0], in[1], axis, s)`. - compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases (last axis, axis 0, a 3-D transformer-shaped gather). This was the last op forcing a fallback in `Bumblebee.Text.question_answering`'s answer-span gather, so the DistilBERT-QA `Nx.Serving` forward now compiles fully native and fused. Wire it into the conformance suite: extend `mode_test` with a `:tag` option (stamp every lane) so the serving test runs its eval / native / native_compiled lanes gated behind `:distilbert_full`.
The opcode wire values are hand-maintained in two places — `Emily.IR`'s `@opcodes` map and the `Opcode` enum + `kOpcodeCount` in `c_src/emily/opcodes.hpp` — with nothing enforcing they agree. A mismatch compiles fine and only misbehaves at runtime. Add `Emily.OpcodeParityTest`, which parses the header and asserts both sides are a unique, gap-free `0..N-1` with N == `kOpcodeCount`. The check is value-based, not name-based, so it doesn't depend on the snake_case names matching the PascalCase enum (`negate`/`Negative`, `fast_rms_norm`/`FastRMSNorm`); a name/value permutation that keeps both contiguous is still caught by the equivalence suite. Expose `Emily.IR.opcodes/0` for the test to read the map.
The example documents are Livebooks (`.livemd`), so name the directory to match. Rename `notebooks/` → `livebooks/` (all 8 files) and update every current-state reference: - mix.exs: the `extras:` paths and the `groups_for_extras` regex, plus the nav-group label `Notebooks:` → `Livebooks:`. - scripts/test-livebooks.sh: `NB_DIR` and the path comment. - README.md / ROADMAP.md: the `livebooks/` path and the HexDocs "Livebooks" section references. CHANGELOG.md and PLAN.md keep their `notebooks/` mentions — those are point-in-time historical records, not current-structure docs.
The window-reduce family (`window_sum`/`window_max`/`window_min`/ `window_product` — average and max pooling) now lowers under `compiler: Emily.Compiler, native: true` instead of raising. Unlike a single MLX op, these are composed (pad -> as_strided sliding view -> reduce), so the cores had to be made shareable first. - op_cores.hpp: hoist `contiguous_strides`/`do_pad`/`sliding_windows_view` out of pooling.cpp's anonymous namespace into `emily::ops`, plus a `WindowReduceKind` enum and `window_reduce_core`. pooling.cpp's eager NIFs now call these, so eager and compiled share one implementation. - opcodes.hpp: add WindowSum/Max/Min/Product (82-85), bump kOpcodeCount, dispatch via `window_reduce_core`. - ir.ex: lower_op clause for the four ops, mirroring `Emily.Backend.apply_window_reduce/6` — operands [input, init], iattrs [window, strides, pad_lo, pad_hi, dilations]. The dtype identity (0/1/±inf or int min/max) is baked as a const operand. - equivalence tests: native-vs-eval bit-identical for 2x2 max/sum/min pooling, padded maxpool, and window_product. A small-CNN forward (conv + maxpool) now lowers fully native under `native_fallback: :raise`, unblocking CNN/avg-pool inference. The maxpool backward (`window_scatter_max`) is still pending for CNN training.
Finish the pooling op family so a small-CNN training step lowers with no fallback: the MaxPool/MinPool backward (`window_scatter_max`/ `window_scatter_min`) and `reverse` (the conv-backward kernel flip). - op_cores.hpp: hoist `window_scatter_core` (the select-and-scatter body) out of pooling.cpp's anonymous namespace into `emily::ops`; the eager NIFs now call it, sharing one implementation with the compiled replay. - opcodes.hpp: add WindowScatterMax/Min (86-87), bump kOpcodeCount to 88, dispatch via `window_scatter_core`. - ir.ex: lower_op for `:window_scatter_max`/`:window_scatter_min` (operands [input, source, init]; iattrs [window, strides, pad_lo, pad_hi]), mirroring `Emily.Backend.apply_window_scatter/7`; and `:reverse`, lowered to one `flip` per axis (reusing the existing flip_core / Flip opcode — no new opcode). - equivalence tests: window_scatter_max/min, grad(window_max) (the maxpool backward path), and reverse — all native-vs-eval bit-identical. A conv + maxpool CNN training step (forward + backward + grad + SGD) now lowers fully native under `native_fallback: :raise`, with a loss bit-identical to the evaluator.
Window reductions with `window_dilations > 1` silently returned wrong
values for windows past the first stride positions, on both the eager
backend and the native Expr compiler (they share `window_reduce_core`).
A dilated kernel axis gets an `as_strided` stride > 1, so the
sliding-window view aliases fewer physical elements than its logical
size (overlapping strides). MLX's reduction then takes its
`GeneralStridedReduce` fast path, which assumes a dense, non-overlapping
layout and reads `product(shape)` contiguous elements — over-running the
aliased buffer. The symptom was data-dependent: large dilated cases
happened to over-read valid in-buffer data, while a tiny `{1, 8}` input
read past the allocation (deterministic 0.0 eager, garbage native).
Materialise the view with `mx::contiguous` before the reduce when any
dilation > 1; the general copy reads element-by-element via the real
strides (always in-bounds), yielding a dense buffer the reducer can
safely fast-path. The common non-dilated pooling path is unchanged and
stays copy-free.
Add regression tests: small-tensor dilated sum/max/min/product vs the
BinaryBackend oracle, and a native-vs-evaluator dilated equivalence case
(restoring the coverage dropped in e32b7cc).
ausimian
changed the base branch from
feat/expr-compiler-livebooks-rename
to
feat/expr-compiler
June 6, 2026 05:37
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.
Lowers the full window/pooling op family so a small-CNN training step
compiles native with no fallback.
feat: lower forward window (pooling) ops—window_sum/_max/_min/_product(average + max pooling). Hoistscontiguous_strides/do_pad/sliding_windows_view+ awindow_reduce_coreintoemily::ops(op_cores.hpp) so the eagerNIFs and the compiled replay share one implementation; opcodes 82–85.
feat: lower window scatter + reverse — CNN training fully native—the MaxPool/MinPool backward (
window_scatter_max/_min, opcodes86–87) +
reverse(the conv-backward kernel flip, lowered to per-axisflip). A conv + maxpool training step (forward + backward + grad +SGD) now lowers fully native, loss bit-identical to the evaluator.
fix: materialize dilated window view before reduce(Dilated window reductions (window_dilations > 1) return wrong values #175) — adilated kernel (
window_dilations > 1) gave an overlappingas_stridedview that MLX's strided-reduce fast path over-read; materialise with
mx::contiguousbefore the reduce. Affected eager + native (sharedcore). Adds dilated regression + equivalence tests.
Touches the Metal/C++ NIF (
op_cores.hpp,opcodes.hpp,pooling.cpp)and
Emily.IR. Equivalence tests are native-vs-evaluator bit-identical.PR 4/5. Stacked on
…-livebooks-rename; lands infeat/expr-compiler.