.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/benchmarks/plot_gemm_benchmark.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_benchmarks_plot_gemm_benchmark.py: Benchmark Gemm: numpy vs onnxruntime vs onnx-light vs onnx-light-cpu ===================================================================== This example compares four ways of computing a ``float32`` general matrix multiplication ``Y = A @ B`` for square matrices of increasing size: * **numpy** - :func:`numpy.matmul` (``A @ B``), which dispatches to the platform BLAS and is used here as the reference baseline. * **onnxruntime** - running a single-node ``Gemm`` ONNX model. * **onnx-light (built-in)** - the same single-node model run through ``onnx-light``'s ``ReferenceEvaluator`` *without* registering ``onnx-light-cpu``, i.e. onnx-light's own (non-SIMD) reference ``Gemm`` kernel. Only measured for the smaller sizes in the grid -- see "Why the built-in curve stops early" below. * **onnx-light + onnx-light-cpu** - the SIMD-accelerated ``Gemm`` kernel that ``onnx-light`` dispatches to after :func:`onnx_light_cpu.register_kernels` installs the optimized kernel implementation. The back-ends compute the same result; the goal is to see how their timings evolve as the matrices grow. Why the built-in curve stops early ----------------------------------- ``onnx-light``'s dispatch table is a single, process-wide table: :func:`onnx_light_cpu.register_kernels` permanently replaces the default domain's ``Gemm`` entry, and a given ``RuntimeSession`` resolves (and caches) its kernels on its *first* run. So the only way to observe the un-accelerated, built-in kernel *and* the accelerated one in the same process is to build a separate model/session for the built-in curve and run it *before* :func:`onnx_light_cpu.register_kernels` is ever called -- which is what this example does. Since the built-in kernel is a plain reference implementation with no SIMD or blocking/packing, its cost grows much faster than the other three back-ends; to keep the benchmark's runtime reasonable it is only measured on the three smallest sizes while the accelerated back-ends continue to 4096. .. GENERATED FROM PYTHON SOURCE LINES 41-46 Setup ----- Report which SIMD level the current CPU provides. The mapping is ``0=None``, ``1=SSE2``, ``2=AVX``, ``3=AVX2`` and ``4=AVX512``. .. GENERATED FROM PYTHON SOURCE LINES 46-101 .. code-block:: Python import argparse import gc import os import time import numpy as np import onnxruntime # ``UNITTEST_GOING=1`` shrinks the benchmark (fewer/smaller sizes) so the example # runs quickly as a unit test while still exercising every code path. unit_test_going = os.environ.get("UNITTEST_GOING", "0") in ("1", "true", "True") parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("-r", "--repeat", type=int, default=10 * (os.cpu_count() or 1)) parser.add_argument("-w", "--warmup", type=int, default=2 * (os.cpu_count() or 1)) parser.add_argument("-t", "--max-repeat-time", type=float, default=1.0) args, _ = parser.parse_known_args() if args.repeat <= 0: parser.error("--repeat must be greater than 0") if args.warmup < 0: parser.error("--warmup must be greater than or equal to 0") if args.max_repeat_time <= 0: parser.error("--max-repeat-time must be greater than 0") # ``onnx-light`` ships ``onnx_light.onnx`` as a drop-in replacement for the # ``onnx`` package; use it to build the model so the example depends on # onnx-light rather than onnx. from onnx_light.onnx import TensorProto, checker, helper from onnx_light.onnx.reference import ReferenceEvaluator from onnx_light_cpu import ( clear_used_kernel_names, detect_simd_level, has_cpu_kernels, register_kernels, registered_kernel_names, SimdLevel, set_kernel_usage_recording, used_kernel_names, ) _SIMD_NAMES = { SimdLevel.NONE: "scalar", SimdLevel.SSE2: "SSE2", SimdLevel.AVX: "AVX", SimdLevel.AVX2: "AVX2", SimdLevel.AVX512: "AVX-512", } assert has_cpu_kernels() level = detect_simd_level() simd_name = _SIMD_NAMES.get(level, level) print(f"CPU kernels available, SIMD level: {level} ({simd_name})") .. rst-class:: sphx-glr-script-out .. code-block:: none CPU kernels available, SIMD level: 4 (AVX-512) .. GENERATED FROM PYTHON SOURCE LINES 102-111 Build the shared ONNX model --------------------------- A single ``Gemm`` node multiplying two 2-D ``float32`` tensors of dynamic shape is enough to benchmark both runtimes. The bias input ``C`` is omitted so the node computes ``A @ B``. ``make_gemm_model`` is called twice: once for the built-in (pre-registration) curve and once for onnxruntime / the accelerated onnx-light-cpu curve, so each gets its own model/graph object (see "Why the built-in curve stops early" above for why that matters). .. GENERATED FROM PYTHON SOURCE LINES 111-160 .. code-block:: Python def make_gemm_model(): graph = helper.make_graph( [helper.make_node("Gemm", ["A", "B"], ["Y"], alpha=1.0, beta=1.0)], "gemm_bench", [ helper.make_tensor_value_info("A", TensorProto.FLOAT, ["M", "K"]), helper.make_tensor_value_info("B", TensorProto.FLOAT, ["K", "N"]), ], [helper.make_tensor_value_info("Y", TensorProto.FLOAT, ["M", "N"])], ) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], ir_version=13) checker.check_model(model) return model model = make_gemm_model() model_bytes = model.SerializeToString() # --------------------------------------------------------------------------- # Timing helper # --------------------------------------------------------------------------- # # Each candidate gets ``--warmup`` untimed calls, then up to ``--repeat`` # measured calls or ``--max-repeat-time`` cumulative seconds. def measure(func, repeat, warmup, max_duration): warmup_duration = 0.0 for _ in range(warmup): start = time.perf_counter() func() warmup_duration += time.perf_counter() - start if warmup_duration >= max_duration: break timings = [] total_duration = 0.0 for _ in range(repeat): start = time.perf_counter() func() duration = time.perf_counter() - start timings.append(duration) total_duration += duration if total_duration >= max_duration: break return float(np.median(timings)) .. GENERATED FROM PYTHON SOURCE LINES 161-168 Prime the built-in onnx-light kernel before registering onnx-light-cpu ---------------------------------------------------------------------- A dedicated model/session is run once while onnx-light-cpu's accelerated ``Gemm`` kernel has not been installed yet, so it resolves and caches onnx-light's own built-in reference kernel. It is timed below alongside the other runtimes on the three smallest sizes. .. GENERATED FROM PYTHON SOURCE LINES 168-208 .. code-block:: Python full_size_grid = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096] size_grid = full_size_grid[:3] if unit_test_going else full_size_grid alone_sizes = full_size_grid[:3] rng = np.random.default_rng(0) alone_model = make_gemm_model() alone_session = ReferenceEvaluator(alone_model) alone_session.run( None, { "A": np.zeros((1, 1), dtype=np.float32), "B": np.zeros((1, 1), dtype=np.float32), }, ) alone_times = {} light_label = "onnx-light + onnx-light-cpu" alone_label = "onnx-light (built-in)" # ``register_kernels()`` needs the ``_cpuregister`` extension, which is only # built with ``ONNX_LIGHT_CPU_WITH_ONNX_LIGHT=ON``. When it is missing (as in # the documentation build) the onnx-light-cpu curve is simply omitted; the # import above stays unconditional. register_kernels() light_session = ReferenceEvaluator(model) def run_light(a, b): return light_session.run(None, {"A": a, "B": b})[0] accelerated_kernel_name = registered_kernel_names()["Gemm"] # onnx-light-cpu kernels record their name on every run (a mutex per call); # only the accelerated curve would pay that cost, so disable recording to keep # the timings fair. Recording is briefly re-enabled below to verify the exact # implementation used for every benchmark size. set_kernel_usage_recording(False) .. GENERATED FROM PYTHON SOURCE LINES 209-215 Run the rest of the benchmark ------------------------------- Each backend runs in its own phase. In particular, the accelerated kernel is measured before NumPy initializes a BLAS pool and before the ONNX Runtime session exists. Inputs are regenerated from the same seed in every phase. .. GENERATED FROM PYTHON SOURCE LINES 215-312 .. code-block:: Python rows_by_size = {size: [size, None, None, None] for size in size_grid} def inputs(): rng = np.random.default_rng(0) for size in size_grid: yield ( size, rng.standard_normal((size, size)).astype(np.float32), rng.standard_normal((size, size)).astype(np.float32), ) for size, a, b in inputs(): set_kernel_usage_recording(True) clear_used_kernel_names() run_light(a, b) accelerated_kernel_names = used_kernel_names() assert accelerated_kernel_name in accelerated_kernel_names, accelerated_kernel_names set_kernel_usage_recording(False) rows_by_size[size][2] = measure( lambda a=a, b=b: run_light(a, b), args.repeat, args.warmup, args.max_repeat_time, ) for size, a, b in inputs(): if size not in alone_sizes: continue alone_times[size] = measure( lambda a=a, b=b: alone_session.run(None, {"A": a, "B": b}), args.repeat, args.warmup, args.max_repeat_time, ) for size, a, b in inputs(): expected = a @ b np.testing.assert_allclose(run_light(a, b), expected, rtol=1e-2, atol=1e-2) if size in alone_sizes: np.testing.assert_allclose( alone_session.run(None, {"A": a, "B": b})[0], expected, rtol=1e-2, atol=1e-2, ) plot_light_results = light_session is not None light_session = None alone_session = None gc.collect() for size, a, b in inputs(): rows_by_size[size][1] = measure( lambda a=a, b=b: a @ b, args.repeat, args.warmup, args.max_repeat_time ) session = onnxruntime.InferenceSession(model_bytes, providers=["CPUExecutionProvider"]) for size, a, b in inputs(): rows_by_size[size][3] = measure( lambda a=a, b=b: session.run(None, {"A": a, "B": b}), args.repeat, args.warmup, args.max_repeat_time, ) for _size, a, b in inputs(): expected = a @ b np.testing.assert_allclose( session.run(None, {"A": a, "B": b})[0], expected, rtol=1e-2, atol=1e-2 ) rows = [tuple(rows_by_size[size]) for size in size_grid] for size, numpy_time, cpu_time, ort_time in rows: alone_text = f"{alone_times[size] * 1e6:10.2f} us" if size in alone_times else "not measured" print( f"size={size:>4}x{size:<4} | numpy={numpy_time * 1e6:10.2f} us | " f"onnx-light={alone_text} | " f"onnx-light-cpu={cpu_time * 1e6:10.2f} us | " f"onnxruntime={ort_time * 1e6:10.2f} us" ) print( "verified onnx-light-cpu Gemm for every benchmark size: " f"accelerated={accelerated_kernel_name}" ) set_kernel_usage_recording(True) sizes = np.array([r[0] for r in rows]) numpy_times = np.array([r[1] for r in rows]) cpu_times = np.array([r[2] for r in rows]) ort_times = np.array([r[3] for r in rows]) alone_grid = np.array(alone_sizes) alone_grid_times = np.array([alone_times[size] for size in alone_sizes]) .. rst-class:: sphx-glr-script-out .. code-block:: none size= 16x16 | numpy= 1.14 us | onnx-light= 3.09 us | onnx-light-cpu= 2.44 us | onnxruntime= 7.46 us size= 32x32 | numpy= 1.89 us | onnx-light= 3.25 us | onnx-light-cpu= 3.05 us | onnxruntime= 8.07 us size= 64x64 | numpy= 5.80 us | onnx-light= 9.20 us | onnx-light-cpu= 8.99 us | onnxruntime= 14.21 us size= 128x128 | numpy= 30.27 us | onnx-light=not measured | onnx-light-cpu= 25.44 us | onnxruntime= 52.44 us size= 256x256 | numpy= 187.62 us | onnx-light=not measured | onnx-light-cpu= 214.88 us | onnxruntime= 345.02 us size= 512x512 | numpy= 1285.39 us | onnx-light=not measured | onnx-light-cpu= 1674.97 us | onnxruntime= 1307.68 us size=1024x1024 | numpy= 10083.61 us | onnx-light=not measured | onnx-light-cpu= 11371.42 us | onnxruntime= 15810.30 us size=2048x2048 | numpy= 79090.71 us | onnx-light=not measured | onnx-light-cpu= 94351.43 us | onnxruntime= 128503.52 us size=4096x4096 | numpy= 621650.11 us | onnx-light=not measured | onnx-light-cpu= 746767.21 us | onnxruntime=1013797.54 us verified onnx-light-cpu Gemm for every benchmark size: accelerated=onnx_light_cpu::Gemm .. GENERATED FROM PYTHON SOURCE LINES 313-325 Plot the timings ---------------- The left panel shows the raw execution time versus the matrix size on a log-log scale. The right panel shows the speed-up relative to **onnxruntime** (the baseline): for each back-end the onnxruntime time is divided by the back-end time, so values above ``1`` are faster than onnxruntime and values below ``1`` are slower. The speed-up is drawn on a logarithmic y-axis so that a given ratio and its reciprocal are equidistant from the ``1`` baseline. The built-in onnx-light curve only has points for the sizes it was measured on (see "Why the built-in curve stops early" above). .. GENERATED FROM PYTHON SOURCE LINES 325-414 .. code-block:: Python import matplotlib.pyplot as plt fig, (ax_time, ax_speedup) = plt.subplots(1, 2, figsize=(11, 4.5)) ax_time.plot(sizes, numpy_times * 1e6, "o--", label="numpy", color="#9b7ec8") ax_time.plot( alone_grid, alone_grid_times * 1e6, "o--", label=alone_label, color="#5cb85c", linewidth=3.5, markersize=9, markerfacecolor="none", markeredgewidth=2, zorder=2, ) if plot_light_results: ax_time.plot( sizes, cpu_times * 1e6, "o-", label=light_label, color="#4a9eff", linewidth=1.5, markersize=4, zorder=3, ) ax_time.plot(sizes, ort_times * 1e6, "o-", label="onnxruntime", color="#f4a259") ax_time.set_xscale("log") ax_time.set_yscale("log") ax_time.set_xlabel("matrix size N (N x N)") ax_time.set_ylabel("time (microseconds)") ax_time.set_title(f"Gemm execution time (SIMD: {simd_name})") ax_time.tick_params(axis="x", labelrotation=20) ax_time.legend() ort_times_by_size = dict(zip(sizes.tolist(), ort_times.tolist(), strict=True)) alone_ort_times = np.array([ort_times_by_size[size] for size in alone_sizes]) ax_speedup.plot(sizes, ort_times / numpy_times, "o--", label="numpy", color="#9b7ec8") ax_speedup.plot( alone_grid, alone_ort_times / alone_grid_times, "o--", label=alone_label, color="#5cb85c", linewidth=3.5, markersize=9, markerfacecolor="none", markeredgewidth=2, zorder=2, ) if plot_light_results: cpu_speedup = ort_times / cpu_times ax_speedup.plot( sizes, cpu_speedup, "o-", label=light_label, color="#4a9eff", linewidth=1.5, markersize=4, zorder=3, ) for size, speedup in zip(sizes, cpu_speedup, strict=True): ax_speedup.annotate( f"{speedup:.2f}x", (size, speedup), xytext=(0, 7), textcoords="offset points", ha="center", fontsize=8, color="#4a9eff", ) ax_speedup.plot(sizes, ort_times / ort_times, "o-", label="onnxruntime", color="#f4a259") ax_speedup.axhline(1.0, color="grey", linewidth=0.8, linestyle=":") ax_speedup.set_xscale("log") ax_speedup.set_yscale("log") ax_speedup.set_xlabel("matrix size N (N x N)") ax_speedup.set_ylabel("speed-up vs onnxruntime") ax_speedup.set_title("Gemm speed-up (onnxruntime = 1)") ax_speedup.tick_params(axis="x", labelrotation=20) ax_speedup.legend() fig.tight_layout() fig.savefig("plot_gemm_benchmark.png") plt.show() .. image-sg:: /auto_examples/benchmarks/images/sphx_glr_plot_gemm_benchmark_001.png :alt: Gemm execution time (SIMD: AVX-512), Gemm speed-up (onnxruntime = 1) :srcset: /auto_examples/benchmarks/images/sphx_glr_plot_gemm_benchmark_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 23.100 seconds) .. _sphx_glr_download_auto_examples_benchmarks_plot_gemm_benchmark.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_gemm_benchmark.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_gemm_benchmark.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_gemm_benchmark.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_