Fix floating point promotion in arithmetic reduction

[Ved235]

Proposed changes

Addresses issue: #4379 and continues on #3909. Modifies routing logic of reduction_op in mlx/backend/cpu/reduce.cpp so that the floating-point accumulator is correctly chosen for arithmetic reduction operations (SumReduce and ProdReduce). As noted in #3909 for apple silicon devices only bfloat16 is being promoted to float32, while float16 and float32 are left unchanged (due to performance implications). However, for devices which use a basic_simd.h bfloat16, float16 are promoted to float32 and float32 is promoted to double.

Benchmark

import time
import mlx.core as mx


N_WARMUP = 5
N_ITERS = 200
ALL_REDUCE_SIZE = 4 * 1024 * 1024
MATRIX_SHAPE = (4096, 4096)


def time_fn(fn):
    for _ in range(N_WARMUP):
        mx.eval(fn())
    mx.synchronize(mx.cpu)

    start = time.perf_counter()
    for _ in range(N_ITERS):
        mx.eval(fn())
    mx.synchronize(mx.cpu)
    return (time.perf_counter() - start) * 1e3 / N_ITERS


def print_table(rows):
    case_width = max(len("case"), *(len(case) for case, _ in rows))
    time_width = len("ms/call")
    print(f"{'case':<{case_width}}  {'ms/call':>{time_width}}")
    print(f"{'-' * case_width}  {'-' * time_width}")
    for case, milliseconds in rows:
        print(f"{case:<{case_width}}  {milliseconds:>{time_width}.5f}")


def benchmark_dtype(name, dtype):
    all_input = mx.ones((ALL_REDUCE_SIZE,), dtype=dtype)
    matrix_input = mx.ones(MATRIX_SHAPE, dtype=dtype)
    mx.eval(all_input, matrix_input)
    mx.synchronize(mx.cpu)

    return [
        (f"{name} sum, all (4M)", time_fn(lambda: mx.sum(all_input))),
        (
            f"{name} sum, axis=-1",
            time_fn(lambda: mx.sum(matrix_input, axis=-1)),
        ),
        (
            f"{name} sum, axis=0 (strided)",
            time_fn(lambda: mx.sum(matrix_input, axis=0)),
        ),
        (
            f"{name} mean, axis=-1",
            time_fn(lambda: mx.mean(matrix_input, axis=-1)),
        ),
    ]



mx.set_default_device(mx.cpu)
rows = []
for name, dtype in [
    ("bf16", mx.bfloat16),
    ("f16", mx.float16),
    ("f32", mx.float32),
]:
    rows.extend(benchmark_dtype(name, dtype))
print_table(rows)

Using this benchmarking script the performance changes are:

Case	Before (ms)	After (ms)	Speedup
bf16 sum, all (4M)	16.89136	2.16033	7.82×
bf16 sum, axis=-1	67.32574	8.10352	8.31×
bf16 sum, axis=0 (strided)	2.22895	0.77700	2.87×
bf16 mean, axis=-1	67.14907	8.05400	8.34×
f16 sum, all (4M)	0.29753	0.29658	1.00×
f16 sum, axis=-1	0.96503	0.96865	1.00×
f16 sum, axis=0 (strided)	0.56356	0.55670	1.01×
f16 mean, axis=-1	0.96479	0.96691	1.00×
f32 sum, all (4M)	0.32213	0.30211	1.07×
f32 sum, axis=-1	1.06188	1.09947	0.97×
f32 sum, axis=0 (strided)	1.07307	1.01758	1.05×
f32 mean, axis=-1	1.07505	1.07238	1.00×

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)