Mojo Benchmarking

Benchmarking measures how fast your code runs so you can identify slow sections and verify that optimizations actually help. Mojo provides a built-in benchmarking module that handles warm-up iterations, timing precision, and statistical noise — giving you reliable numbers instead of guesses.

Why Formal Benchmarking Matters

Naive timing:
  start = time.now()
  do_work()
  end = time.now()
  print(end - start)   ← unreliable: OS scheduling, CPU frequency scaling,
                          cold caches, branch predictor state all affect this

Proper benchmarking:
  ✓ Runs the function many times (warm-up + measurement iterations)
  ✓ Disables compiler optimizations that would remove the work
  ✓ Reports mean, min, max, and standard deviation
  ✓ Accounts for system noise

The Benchmark Module

from benchmark import Bench, BenchConfig, Bencher, keep

fn square_sum(n: Int) -> Int:
    var total = 0
    for i in range(n):
        total += i * i
    return total

fn bench_square_sum(inout b: Bencher) raises:
    @always_inline
    @parameter
    fn call_fn():
        var result = square_sum(1000)
        keep(result)   # prevents compiler from optimizing the call away

    b.iter[call_fn]()

fn main() raises:
    var bench = Bench(BenchConfig(num_repetitions=5))
    bench.run[bench_square_sum]("square_sum_1000")
    bench.dump_report()

Output (approximate):

-------------------------- Benchmark Report --------------------------
Name              | Avg (ns) | Min (ns) | Max (ns) | Iters
------------------|----------|----------|----------|------
square_sum_1000   |   1234   |   1198   |   1289   | 5000

The keep() Function

The Mojo compiler is aggressive about dead-code elimination. If it sees that the result of a function is never used, it may remove the function call entirely — making your benchmark measure nothing. keep(result) prevents this by signaling that the result matters.

# Wrong — compiler may remove square_sum entirely:
fn bench_bad(inout b: Bencher) raises:
    @parameter
    fn call_fn():
        _ = square_sum(1000)   # result discarded → may be optimized away
    b.iter[call_fn]()

# Correct — keep() forces the computation to happen:
fn bench_good(inout b: Bencher) raises:
    @parameter
    fn call_fn():
        var result = square_sum(1000)
        keep(result)           # result is "used" → not removed
    b.iter[call_fn]()

Comparing Two Implementations

from benchmark import Bench, BenchConfig, Bencher, keep
from memory import UnsafePointer

fn scalar_sum(data: UnsafePointer[Float32], n: Int) -> Float32:
    var total: Float32 = 0.0
    for i in range(n):
        total += data[i]
    return total

fn simd_sum(data: UnsafePointer[Float32], n: Int) -> Float32:
    var acc = SIMD[DType.float32, 8].splat(0.0)
    var i = 0
    while i + 8 <= n:
        acc += SIMD[DType.float32, 8].load(data + i)
        i += 8
    while i < n:
        acc[0] += data[i]
        i += 1
    return acc.reduce_add()

fn main() raises:
    let n = 10_000
    var data = UnsafePointer[Float32].alloc(n)
    for i in range(n):
        data.init_pointee_copy(1.0)

    var bench = Bench(BenchConfig())

    fn bench_scalar(inout b: Bencher) raises:
        @parameter
        fn run():
            keep(scalar_sum(data, n))
        b.iter[run]()

    fn bench_simd(inout b: Bencher) raises:
        @parameter
        fn run():
            keep(simd_sum(data, n))
        b.iter[run]()

    bench.run[bench_scalar]("scalar")
    bench.run[bench_simd]("simd")
    bench.dump_report()

    for i in range(n):
        (data + i).destroy_pointee()
    data.free()

Reading Benchmark Output

Name     | Avg (ns) | Speedup
---------|----------|--------
scalar   |   12400  |  1.0×
simd     |    1600  |  7.75×

Interpreting:
  Avg = average time per single call in nanoseconds
  Lower = faster

  simd is 7.75× faster than scalar on this data size.
  This matches the expected ~8× from using 8-wide SIMD.

BenchConfig Options

var config = BenchConfig(
    num_repetitions = 10,      # outer repetitions for statistics
    warmup_iters = 100,        # iterations before timing starts
    max_iters = 10_000,        # max timing iterations per repetition
    min_runtime_secs = 0.5,    # run for at least this many seconds
)

Profiling Tips

Benchmark at Representative Scale

A function that sums 10 elements may be dominated by call overhead. Benchmark at the scale you actually use in production — typically thousands to millions of elements for numerical code.

Benchmark in Release Mode

# Debug builds include extra checks and run slower.
# Always benchmark the optimized build:
magic run mojo -O3 my_benchmark.mojo

Isolate One Variable at a Time

Change one thing between benchmark runs — tile size, SIMD width, or algorithm choice. Changing multiple things at once makes it impossible to determine what caused the speedup.

Key Takeaways

Mojo's benchmark module provides statistically sound timing with warm-up iterations and multiple repetitions. Always use keep() to prevent the compiler from removing the computation being measured. Compare implementations head-to-head in the same benchmark run to eliminate machine state differences. Benchmark at production-scale data sizes. Change one variable per experiment so results are interpretable. Benchmark numbers guide optimization — always measure before and after an optimization to confirm it actually helps.

Leave a Comment

Your email address will not be published. Required fields are marked *