lablup / lablup/mlxcel

fix(mlx-cpp): stop casting comparison inputs to bool in compiled kernels

Open
#1,792 0 comments 0 reactions 0 assignees View on GitHub
area:core priority:medium status:ready type:bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.