Benchmark Abs: onnxruntime vs onnx-light#

This example compares the built-in Abs kernel in onnx-light with onnxruntime and numpy.abs() for vectors ranging from one hundred to one hundred million elements. The comparison is repeated for three element types: float32, float16 and bfloat16. The bfloat16 values rely on the ml_dtypes NumPy extension, which provides a native bfloat16 dtype.

onnxruntime only ships a CPU Abs implementation for float32 and float16; there is no bfloat16 kernel. The benchmark therefore skips the onnxruntime measurement and speed-up plot for bfloat16.

The execution benchmark warms each runtime and reports median durations. onnx-light is measured before the ONNX Runtime session is constructed: keeping both persistent CPU pools in one process while alternating calls causes one runtime’s spinning workers to perturb the other runtime’s measurement. The two onnx-light series use the same prepared evaluator with either a NumPy array or a pre-built runtime Tensor as input.

from __future__ import annotations

import os
import time

import matplotlib.pyplot
import ml_dtypes
import numpy
import onnxruntime
import onnx_light.onnx.helper as oh
from onnx_light.onnx import TensorProto, checker
from onnx_light.onnx.reference import ReferenceEvaluator
from onnx_light.onnx_py import _onnxpykernels

runtime = _onnxpykernels.runtime
ORT_MAX_IR_VERSION = 13

Element types under test#

Each entry holds the label used in the report, the ONNX TensorProto type, the matching NumPy dtype and whether onnxruntime provides a CPU Abs kernel for that type. bfloat16 is materialized through ml_dtypes; onnxruntime has no bfloat16 Abs kernel, so it is excluded from that comparison.

DTYPES = [
    ("float32", TensorProto.FLOAT, numpy.dtype(numpy.float32), True),
    ("float16", TensorProto.FLOAT16, numpy.dtype(numpy.float16), True),
    ("bfloat16", TensorProto.BFLOAT16, numpy.dtype(ml_dtypes.bfloat16), False),
]


def make_abs_model(elem_type: int):
    """Creates a dynamic one-dimensional Abs model for a given element type."""

    graph = oh.make_graph(
        [oh.make_node("Abs", ["X"], ["Y"])],
        "abs_benchmark",
        [oh.make_tensor_value_info("X", elem_type, ["N"])],
        [oh.make_tensor_value_info("Y", elem_type, ["N"])],
    )
    model = oh.make_model(graph, opset_imports=[oh.make_opsetid("", 18)])
    model.ir_version = min(model.ir_version, ORT_MAX_IR_VERSION)
    checker.check_model(model)
    return model


def measure(function, repeat: int, warmup: int = 3, number: int = 1) -> float:
    """Measures a callable after warm-up and returns its median time per call."""

    for _ in range(warmup):
        function()
    timings = []
    for _ in range(repeat):
        start = time.perf_counter()
        for _ in range(number):
            function()
        timings.append((time.perf_counter() - start) / number)
    return float(numpy.median(timings))

Measurement grid#

Normal execution uses the complete logarithmic grid. Documentation tests use two small vectors to keep the gallery build fast.

if os.environ.get("UNITTEST_GOING") == "1":
    size_grid = [100, 1_000]
    minimum_repeat = 3
    warmup = 1
else:
    size_grid = [10**power for power in range(2, 9)]
    minimum_repeat = 7
    warmup = 3


def benchmark_dtype(label: str, elem_type: int, np_dtype, ort_supported: bool) -> dict:
    """Benchmarks the Abs kernel for a single element type.

    Both runtimes receive inputs generated from the same seed and are warmed
    before timing. All ``onnx-light`` measurements finish before the ONNX
    Runtime session is constructed, keeping their persistent CPU pools from
    perturbing each other. ``onnxruntime`` is only exercised when it provides a
    CPU ``Abs`` kernel for ``elem_type``.

    Returns:
        A mapping with the measured ``sizes`` and, for each backend, the
        median execution times. ``onnxruntime`` times are ``None`` when the
        element type is unsupported.
    """

    model = make_abs_model(elem_type)
    model_bytes = model.SerializeToString()

    onnx_light_session = ReferenceEvaluator(model)

    def run_onnx_light(values):
        """Runs the built-in onnx-light Abs kernel."""

        return onnx_light_session.run(None, {"X": values})[0]

    def run_onnx_light_tensor(tensor):
        """Runs onnx-light with a pre-built runtime Tensor."""

        return onnx_light_session.run(None, {"X": tensor})[0]

    def make_input_tensor(values):
        """Creates a zero-copy runtime Tensor over a NumPy input."""

        return runtime.tensor_from_numpy(
            "X", int(elem_type), list(values.shape), values.view(numpy.uint8), copy=False
        )

    def run_onnxruntime(values):
        """Runs the ONNX Runtime Abs kernel."""

        return ort_session.run(None, {"X": values})[0]

    random_generator = numpy.random.default_rng(0)
    rows_by_size = {}
    for size in size_grid:
        values = random_generator.uniform(-100.0, 100.0, size=size).astype(np_dtype)
        expected = numpy.abs(values)
        repeat = max(minimum_repeat, min(200, 2_000_000 // size))
        number = max(1, min(20, 10_000_000 // size))

        numpy_time = measure(lambda values=values: numpy.abs(values), repeat, warmup, number)
        onnx_light_time = measure(
            lambda values=values: run_onnx_light(values), repeat, warmup, number
        )
        input_tensor = make_input_tensor(values)
        onnx_light_tensor_time = measure(
            lambda tensor=input_tensor: run_onnx_light_tensor(tensor), repeat, warmup, number
        )
        numpy.testing.assert_array_equal(run_onnx_light(values), expected)
        numpy.testing.assert_array_equal(run_onnx_light_tensor(input_tensor), expected)
        rows_by_size[size] = [size, numpy_time, onnx_light_time, onnx_light_tensor_time, None]

    if ort_supported:
        ort_session = onnxruntime.InferenceSession(
            model_bytes, providers=["CPUExecutionProvider"]
        )
        random_generator = numpy.random.default_rng(0)
        for size in size_grid:
            values = random_generator.uniform(-100.0, 100.0, size=size).astype(np_dtype)
            expected = numpy.abs(values)
            repeat = max(minimum_repeat, min(200, 2_000_000 // size))
            number = max(1, min(20, 10_000_000 // size))
            ort_time = measure(
                lambda values=values: run_onnxruntime(values), repeat, warmup, number
            )
            numpy.testing.assert_array_equal(run_onnxruntime(values), expected)
            rows_by_size[size][4] = ort_time

    rows = [tuple(rows_by_size[size]) for size in size_grid]
    for size, numpy_time, onnx_light_time, onnx_light_tensor_time, ort_time in rows:
        ort_report = "n/a" if ort_time is None else f"{ort_time * 1e6:10.2f} us"
        ratio_report = "n/a" if ort_time is None else f"{onnx_light_time / ort_time:5.2f}x"
        print(
            f"[{label:>8}] size={size:>9} | numpy={numpy_time * 1e6:10.2f} us | "
            f"onnx-light={onnx_light_time * 1e6:10.2f} us | "
            f"onnx-light (Tensor)={onnx_light_tensor_time * 1e6:10.2f} us | "
            f"onnxruntime={ort_report} | onnx-light / onnxruntime={ratio_report}"
        )

    return {
        "label": label,
        "ort_supported": ort_supported,
        "sizes": numpy.array([row[0] for row in rows]),
        "numpy_times": numpy.array([row[1] for row in rows]),
        "onnx_light_times": numpy.array([row[2] for row in rows]),
        "onnx_light_tensor_times": numpy.array([row[3] for row in rows]),
        "ort_times": None if not ort_supported else numpy.array([row[4] for row in rows]),
    }

Measure steady-state execution for every element type#

results = [benchmark_dtype(*entry) for entry in DTYPES]

[ float32] size=      100 | numpy=      0.33 us | onnx-light=      1.68 us | onnx-light (Tensor)=      1.66 us | onnxruntime=      4.17 us | onnx-light / onnxruntime= 0.40x
[ float32] size=     1000 | numpy=      0.50 us | onnx-light=      1.75 us | onnx-light (Tensor)=      1.73 us | onnxruntime=      4.36 us | onnx-light / onnxruntime= 0.40x
[ float32] size=    10000 | numpy=      0.91 us | onnx-light=      2.31 us | onnx-light (Tensor)=      2.31 us | onnxruntime=      3.26 us | onnx-light / onnxruntime= 0.71x
[ float32] size=   100000 | numpy=      6.38 us | onnx-light=      5.51 us | onnx-light (Tensor)=      5.30 us | onnxruntime=     11.22 us | onnx-light / onnxruntime= 0.49x
[ float32] size=  1000000 | numpy=     74.49 us | onnx-light=     46.27 us | onnx-light (Tensor)=     42.24 us | onnxruntime=     86.28 us | onnx-light / onnxruntime= 0.54x
[ float32] size= 10000000 | numpy=   3636.01 us | onnx-light=   1298.69 us | onnx-light (Tensor)=   1315.73 us | onnxruntime=   1676.94 us | onnx-light / onnxruntime= 0.77x
[ float32] size=100000000 | numpy=  32558.93 us | onnx-light=  17522.05 us | onnx-light (Tensor)=  17665.97 us | onnxruntime=  18184.28 us | onnx-light / onnxruntime= 0.96x
[ float16] size=      100 | numpy=      0.30 us | onnx-light=      1.65 us | onnx-light (Tensor)=      1.62 us | onnxruntime=      3.39 us | onnx-light / onnxruntime= 0.49x
[ float16] size=     1000 | numpy=      0.44 us | onnx-light=      1.65 us | onnx-light (Tensor)=      1.63 us | onnxruntime=      3.31 us | onnx-light / onnxruntime= 0.50x
[ float16] size=    10000 | numpy=      0.62 us | onnx-light=      1.90 us | onnx-light (Tensor)=      1.86 us | onnxruntime=      4.45 us | onnx-light / onnxruntime= 0.43x
[ float16] size=   100000 | numpy=      2.29 us | onnx-light=      3.73 us | onnx-light (Tensor)=      3.74 us | onnxruntime=     21.74 us | onnx-light / onnxruntime= 0.17x
[ float16] size=  1000000 | numpy=     37.07 us | onnx-light=     28.60 us | onnx-light (Tensor)=     26.14 us | onnxruntime=    107.52 us | onnx-light / onnxruntime= 0.27x
[ float16] size= 10000000 | numpy=    669.29 us | onnx-light=    404.92 us | onnx-light (Tensor)=    372.82 us | onnxruntime=   3584.86 us | onnx-light / onnxruntime= 0.11x
[ float16] size=100000000 | numpy=  15656.69 us | onnx-light=   8865.22 us | onnx-light (Tensor)=   9422.36 us | onnxruntime=  44869.96 us | onnx-light / onnxruntime= 0.20x
[bfloat16] size=      100 | numpy=      0.31 us | onnx-light=      2.87 us | onnx-light (Tensor)=      1.73 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a
[bfloat16] size=     1000 | numpy=      0.72 us | onnx-light=      2.89 us | onnx-light (Tensor)=      1.76 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a
[bfloat16] size=    10000 | numpy=      3.40 us | onnx-light=      3.23 us | onnx-light (Tensor)=      2.09 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a
[bfloat16] size=   100000 | numpy=     30.26 us | onnx-light=      5.28 us | onnx-light (Tensor)=      4.13 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a
[bfloat16] size=  1000000 | numpy=    301.28 us | onnx-light=     31.21 us | onnx-light (Tensor)=     25.14 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a
[bfloat16] size= 10000000 | numpy=   3026.99 us | onnx-light=    410.80 us | onnx-light (Tensor)=    358.44 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a
[bfloat16] size=100000000 | numpy=  40155.46 us | onnx-light=   9074.71 us | onnx-light (Tensor)=   8995.28 us | onnxruntime=n/a | onnx-light / onnxruntime=n/a

Plot execution time and relative speed#

One row is drawn per element type. The left panel shows raw inference time. The right panel shows the speed-up relative to onnxruntime when it provides a CPU kernel, on a logarithmic scale so that speed-ups and slowdowns are equally readable around the baseline. No speed-up is reported for bfloat16 because changing the baseline to NumPy would make that row incomparable.

figure, axes = matplotlib.pyplot.subplots(
    len(results), 2, figsize=(12, 4.5 * len(results)), squeeze=False
)

for row_index, result in enumerate(results):
    label = result["label"]
    sizes = result["sizes"]
    numpy_times = result["numpy_times"]
    onnx_light_times = result["onnx_light_times"]
    onnx_light_tensor_times = result["onnx_light_tensor_times"]
    ort_times = result["ort_times"]

    time_axis = axes[row_index][0]
    speedup_axis = axes[row_index][1]

    time_axis.plot(sizes, numpy_times * 1e6, "o--", label="numpy", color="#9b7ec8")
    time_axis.plot(sizes, onnx_light_times * 1e6, "o-", label="onnx-light", color="#5cb85c")
    time_axis.plot(
        sizes, onnx_light_tensor_times * 1e6, "s:", label="onnx-light (Tensor)", color="#1b5e20"
    )
    if ort_times is not None:
        time_axis.plot(sizes, ort_times * 1e6, "o-", label="onnxruntime", color="#f4a259")
    time_axis.set_xscale("log")
    time_axis.set_yscale("log")
    time_axis.set_xlabel("array size (elements)")
    time_axis.set_ylabel("time (microseconds)")
    time_axis.set_title(f"Abs execution time ({label})")
    time_axis.legend()

    if ort_times is None:
        speedup_axis.axis("off")
        speedup_axis.text(
            0.5,
            0.5,
            "No ONNX Runtime CPU bfloat16 Abs kernel\nspeed-up not reported",
            ha="center",
            va="center",
            transform=speedup_axis.transAxes,
        )
        continue

    # The baseline itself is a flat line at 1.0, shown by the reference
    # ``axhline`` below, so it is not plotted as its own series.
    speedup_axis.plot(sizes, ort_times / numpy_times, "o--", label="numpy", color="#9b7ec8")
    onnx_light_speedups = ort_times / onnx_light_times
    speedup_axis.plot(sizes, onnx_light_speedups, "o-", label="onnx-light", color="#5cb85c")
    for size, speedup in zip(sizes, onnx_light_speedups, strict=True):
        speedup_axis.annotate(
            f"{speedup:.2f}x",
            (size, speedup),
            xytext=(0, 6),
            textcoords="offset points",
            ha="center",
            fontsize=7,
            color="#3d803d",
        )
    speedup_axis.plot(
        sizes,
        ort_times / onnx_light_tensor_times,
        "s:",
        label="onnx-light (Tensor)",
        color="#1b5e20",
    )
    speedup_axis.axhline(
        1.0, color="grey", linewidth=0.8, linestyle=":", label="onnxruntime (baseline)"
    )
    speedup_axis.set_xscale("log")
    speedup_axis.set_yscale("log")
    speedup_axis.set_xlabel("array size (elements)")
    speedup_axis.set_ylabel("speed-up vs onnxruntime")
    speedup_axis.set_title(f"Abs speed-up ({label}, onnxruntime = 1)")
    speedup_axis.legend()

Abs execution time (float32), Abs speed-up (float32, onnxruntime = 1), Abs execution time (float16), Abs speed-up (float16, onnxruntime = 1), Abs execution time (bfloat16)

onnx-light is expected to beat onnxruntime on the smallest float32 vector, where the fixed per-call overhead dominates.

float32_result = results[0]
float32_speedups = float32_result["ort_times"] / float32_result["onnx_light_times"]
assert float32_speedups[0] > 1.0, (
    "onnx-light is expected to be faster than onnxruntime for the first (smallest) size, "
    f"got a speed-up of {float32_speedups[0]:.2f}x for size {float32_result['sizes'][0]}"
)

figure.tight_layout()
figure.savefig("plot_abs_benchmark.png")