Note
Go to the end to download the full example code.
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.
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
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"
)
Total running time of the script: (0 minutes 0.625 seconds)