.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/benchmarks/plot_gather_block_quantized_benchmark.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_benchmarks_plot_gather_block_quantized_benchmark.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_benchmarks_plot_gather_block_quantized_benchmark.py:


Packed embedding gather against ONNX Runtime
============================================

Compare ``com.microsoft::GatherBlockQuantized`` on packed 2-, 4-, and 8-bit
embedding tables. The C++ backend fixtures retain packed weights: neither
runtime receives an expanded floating-point checkpoint. Check numerical parity
before timing, and feed backend tensors directly to onnx-light-cpu.

.. GENERATED FROM PYTHON SOURCE LINES 10-69




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    test_cpu_gather_block_quantized_embedding4096x256_ids16_bits2_scalesfp32_uint8_benchmark: onnx-light-cpu=10.8 us, ORT=10.1 us
    test_cpu_gather_block_quantized_embedding4096x256_ids16_bits4_scalesfp16_uint8_benchmark: onnx-light-cpu=14.1 us, ORT=9.3 us
    test_cpu_gather_block_quantized_embedding4096x256_ids16_bits8_scalesfp32_uint8_benchmark: onnx-light-cpu=10.1 us, ORT=7.3 us






|

.. code-block:: Python


    import os
    import statistics
    import time

    import numpy as np
    import onnxruntime

    from onnx_light.onnx.backend import TestMode, collect_test_cases_by_name
    from onnx_light.onnx.helper import tensor_dtype_to_np_dtype
    from onnx_light.onnx.reference import ReferenceEvaluator
    from onnx_light_cpu import register_backend_test_cases, register_kernels

    register_kernels()
    register_backend_test_cases()


    def to_numpy(tensor):
        """Views a backend fixture as an ORT input."""
        return np.frombuffer(
            tensor.raw_data(), dtype=tensor_dtype_to_np_dtype(int(tensor.data_type))
        ).reshape(tuple(tensor.shape))


    def measure(session, feeds, repeat):
        """Returns the median of a fixed, bounded number of warmed calls."""
        session.run(None, feeds)
        samples = []
        for _ in range(repeat):
            begin = time.perf_counter()
            session.run(None, feeds)
            samples.append(time.perf_counter() - begin)
        return statistics.median(samples)


    repeat = 2 if os.environ.get("UNITTEST_GOING") else 30
    cases = collect_test_cases_by_name(
        "^test_cpu_gather_block_quantized_.*_benchmark$",
        mode=TestMode.BENCHMARK,
        generate_benchmark_expected_outputs=False,
    )
    assert cases, "GatherBlockQuantized backend benchmark fixtures must be built."
    for case in cases:
        names = [value.name for value in case.model.graph.input]
        cpu_feeds = dict(zip(names, case.data_sets[0].inputs, strict=True))
        ort_feeds = {name: to_numpy(value) for name, value in cpu_feeds.items()}
        cpu = ReferenceEvaluator(
            case.model, cpu_execution={"num_threads": 1, "affinity_policy": "none"}
        )
        options = onnxruntime.SessionOptions()
        options.intra_op_num_threads = 1
        ort = onnxruntime.InferenceSession(
            case.model.SerializeToString(), options, providers=["CPUExecutionProvider"]
        )
        np.testing.assert_array_equal(cpu.run(None, cpu_feeds)[0], ort.run(None, ort_feeds)[0])
        print(
            f"{case.name}: onnx-light-cpu={measure(cpu, cpu_feeds, repeat) * 1e6:.1f} us, "
            f"ORT={measure(ort, ort_feeds, repeat) * 1e6:.1f} us"
        )


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 0.625 seconds)


.. _sphx_glr_download_auto_examples_benchmarks_plot_gather_block_quantized_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_gather_block_quantized_benchmark.ipynb <plot_gather_block_quantized_benchmark.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: plot_gather_block_quantized_benchmark.py <plot_gather_block_quantized_benchmark.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: plot_gather_block_quantized_benchmark.zip <plot_gather_block_quantized_benchmark.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_