fix(mlx-cpp): stop casting comparison inputs to bool in compiled kernels
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Problem
The mlxcel overlay of MLX's Metal compiled-kernel generator casts each input whose type differs from the op's output type to the output type. It exists for mixed bf16/f32 arithmetic (commit 1923da32: `Divide(bfloat16_t, float)` failed template deduction on macOS 26.4). A comparison (`Equal`, `NotEqual`, `Less`, `LessEqual`, `Greater`, `GreaterEqual`) outputs `bool`, so every non-bool input is cast to `bool` first, same-dtype inputs included: `Less(0.2, 0.7)` becomes `Less(true, true)`, which is `false`. `isnan` and `isinf` lower to `NotEqual` / `Equal` and break too. The result is silently wrong, never an error. CUDA has no `compiled.cpp` overlay (none under `patches/mlx/backend/cuda/` or `patches-cuda/`) and is unaffected.
Latent today: no compiled function in the tree contains a comparison (every `mlx::core::compile` and `compile_shapeless_audited` site in `mlx_cxx_bridge.cpp` and `mlx_cxx_kernels.cpp` checked), and the #1392 audit found all 20 shapeless sites equal to eager on Metal. The compiled min-p filter removed in PR #1391 (#1379) had this shape and returned its input unfiltered on Metal, mechanism undiagnosed (`src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:5483-5493`); this cast reproduces that symptom. Found during review of PR #1772.
## Evidence
- `src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp:2` (overlay header) and `:220-244` (the cast keyed on `out_type`, which is the overlay's whole delta against upstream 81ba1c6a).
- Probe against the in-tree `libmlx.a` (81ba1c6a, M1 Ultra), compiled vs eager: `where(less(a, b), a, b)` with f32 `a = [0.2, -0.5, 0.7, 0, 3]`, `b = [0.7, 0.3, 0.2, 1, 3]` gives `[0.7, 0.3, 0.2, 0, 3]` instead of `[0.2, -0.5, 0.2, 0, 3]`; `greater(2a, b)` is all false; an f16/f32 comparison diverges; `where(isnan(x), 7, x)` passes NaN through; the min-p graph (`softmax`, `max * 0.3`, `greater_equal`, `where`) returns its input unfiltered. A bf16/f32 arithmetic chain still matches, and upstream MLX 0.32.2 matches eager on every graph.
## Proposed fix
Exclude bool-output primitives from the output-type cast, or cast mismatched inputs to the promoted type of the op's inputs instead of the output type. Either keeps the arithmetic case the cast exists for. Keep the overlay delta small, since every MLX pin bump three-way merges it.
## Acceptance criteria
- [ ] Compiled comparisons on same-dtype f32 inputs and on mixed f16/f32 inputs, and a compiled `isnan`, match the uncompiled result.
- [ ] The mixed bf16/f32 arithmetic case still compiles and matches eager.
## Verification
A new `mlxcel-core` test that fails before the fix (through a test-only bridge entry, since no generic compile call is exposed to Rust), then `cargo test --workspace --profile test-fast --features metal,accelerate` on a Metal host.
Contributor guide
Research direction
Start with src/lib/mlx-cpp/patches/mlx/backend/metal/compiled.cpp:220-244 and inspect how the out_type cast affects comparison primitives. Add a test-only bridge entry for compiled comparisons, then run the mlxcel-core test and cargo test --workspace --profile test-fast --features metal,accelerate on a Metal host. Done means same- and mixed-dtype comparisons and isnan match eager while bf16/f32 arithmetic remains correct.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- compilers, machine-learning, testing
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100