Skip to content

Lower indexed_put / indexed_add (scatter) in the native Expr compiler - #184

Merged
ausimian merged 1 commit into
feat/expr-compilerfrom
feat/expr-compiler-scatter
Jun 6, 2026
Merged

Lower indexed_put / indexed_add (scatter) in the native Expr compiler#184
ausimian merged 1 commit into
feat/expr-compilerfrom
feat/expr-compiler-scatter

Conversation

@ausimian

@ausimian ausimian commented Jun 6, 2026

Copy link
Copy Markdown
Owner

Nx.indexed_put and Nx.indexed_add (scatter) now lower under
compiler: Emily.Compiler, native: true instead of forcing a graceful
fallback. Same shape as the FFT / take_along_axis additions — the MLX
ops and eager NIFs already existed (Native.{scatter,scatter_add}
mx::{scatter,scatter_add}); only the compiler path was missing.
Implementing the one shared scatter path closes both ops.

What surfaced this

Continuing to run the Whisper speech_to_text serving with native: true
logged a fallback on :indexed_put. With this change that path compiles
native for the common (MLX-scatter-compatible) index layout.

Changes

  • c_src/emily/opcodes.hpp — add Scatter/ScatterAdd (92–93), bump
    kOpcodeCount to 94, dispatch to mx::scatter (overwrite) /
    mx::scatter_add (accumulate). operands [target, updates, idx0, ...];
    iattrs [[axes...]].
  • lib/emily/ir.ex — add the opcodes; one lower_op clause for
    :indexed_put/:indexed_add, mirroring Emily.Backend's apply_scatter:
    reuse the existing gather index-split + scatter_gather_compatible?
    helpers, port updates_shape_for_scatter, reshape updates into MLX's
    layout, emit the scatter opcode, coerce to out.type.
  • compiler_equivalence_test.exs — native-vs-evaluator bit-identical
    cases: 2-D grid put/add, duplicate-index accumulation, partial-axis
    (axes: [0]) whole-row writes.

Behaviour note

Like the native gather, only MLX-scatter-compatible index layouts
(trailing dim == #axes, ≥1 batch dim) lower; other layouts raise during
lowering → the evaluator handles them under native_fallback: :eval
(unchanged from today, where any indexed_put falls back). Strict
improvement for the common layout.

Verification

  • New scatter cases pass under native: true (suite default
    native_fallback: :raise, so a fallback would raise — passing proves
    full native lowering). Opcode-parity test confirms the IR↔C++ lockstep
    (92–93 / kOpcodeCount 94).
  • mix precommit green: compile (warnings-as-errors), format, credo
    --strict, docs, 689 tests, 0 failures.

Nx.indexed_put and Nx.indexed_add (scatter) now lower under
`compiler: Emily.Compiler, native: true` for MLX-scatter-compatible index
layouts, instead of forcing a graceful fallback to the evaluator. The MLX
ops and eager NIFs already existed (Native.{scatter,scatter_add} ->
mx::{scatter,scatter_add}); only the compiler path was missing.

- opcodes.hpp: add Scatter/ScatterAdd (92-93), bump kOpcodeCount to 94,
  dispatch to mx::scatter (overwrite) / mx::scatter_add (accumulate).
  operands [target, updates, idx0, ...]; iattrs [[axes...]].
- ir.ex: add the opcodes; lower :indexed_put/:indexed_add (one shared
  clause), mirroring Emily.Backend's apply_scatter — reuse the existing
  gather index-split + scatter_gather_compatible? helpers, port
  updates_shape_for_scatter. Incompatible index layouts raise (no
  fallback), matching the native gather; the evaluator handles them under
  native_fallback: :eval.
- compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases
  for indexed_put/indexed_add (2-D grid, duplicate-index accumulation,
  partial-axis whole-row writes).

Another op surfaced while running the native compiler on the Whisper
livebook.
@ausimian
ausimian merged commit d02b252 into feat/expr-compiler Jun 6, 2026
3 checks passed
@ausimian
ausimian deleted the feat/expr-compiler-scatter branch June 6, 2026 07:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant