benchmarking

Synchronized JAX compilation, execution, and memory measurements.

benchmark_jax_function(function, *args, iterations=25, warmup=1, clear_caches=True, include_hlo=False, compare_precision=False, jit_kwargs=None, call_kwargs=None)

Collect synchronized JAX timing, memory, HLO, cost, and precision stats.

Source code in jaxquantum/utils/benchmarking.py

def benchmark_jax_function(
    function: Callable,
    *args,
    iterations: int = 25,
    warmup: int = 1,
    clear_caches: bool = True,
    include_hlo: bool = False,
    compare_precision: bool = False,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Collect synchronized JAX timing, memory, HLO, cost, and precision stats."""
    if compare_precision:
        return benchmark_precision(
            function,
            *args,
            iterations=iterations,
            warmup=warmup,
            clear_caches=clear_caches,
            include_hlo=include_hlo,
            jit_kwargs=jit_kwargs,
            call_kwargs=call_kwargs,
        )
    if iterations < 1 or warmup < 0:
        raise ValueError("iterations must be positive and warmup non-negative")
    return _benchmark_once(
        function,
        args,
        iterations,
        warmup,
        clear_caches,
        include_hlo,
        jit_kwargs,
        dict(call_kwargs or {}),
    )[0]

benchmark_precision(function, *args, iterations=25, warmup=1, clear_caches=True, include_hlo=False, jit_kwargs=None, call_kwargs=None)

Compare double and single precision performance and output accuracy.

Source code in jaxquantum/utils/benchmarking.py

def benchmark_precision(
    function: Callable,
    *args,
    iterations: int = 25,
    warmup: int = 1,
    clear_caches: bool = True,
    include_hlo: bool = False,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Compare double and single precision performance and output accuracy."""
    if iterations < 1 or warmup < 0:
        raise ValueError("iterations must be positive and warmup non-negative")
    original_x64 = jax.config.x64_enabled
    reports = {}
    outputs = {}
    try:
        for name, enabled, real_dtype, complex_dtype in (
            ("double", True, np.float64, np.complex128),
            ("single", False, np.float32, np.complex64),
        ):
            jax.config.update("jax_enable_x64", enabled)
            precision_args = _cast_precision(args, real_dtype, complex_dtype)
            precision_kwargs = _cast_precision(
                dict(call_kwargs or {}),
                real_dtype,
                complex_dtype,
            )
            reports[name], output = _benchmark_once(
                function,
                precision_args,
                iterations,
                warmup,
                clear_caches,
                include_hlo,
                jit_kwargs,
                precision_kwargs,
            )
            outputs[name] = jax.device_get(output)
            del output
    finally:
        jax.config.update("jax_enable_x64", original_x64)
        if clear_caches:
            jax.clear_caches()

    double_timing = reports["double"]["timings_s"]
    single_timing = reports["single"]["timings_s"]
    double_memory = reports["double"]["memory_bytes"]
    single_memory = reports["single"]["memory_bytes"]
    return {
        "double": reports["double"],
        "single": reports["single"],
        "accuracy": _accuracy_stats(outputs["double"], outputs["single"]),
        "single_vs_double": {
            "cold_speedup": _ratio(
                double_timing["cold_total"],
                single_timing["cold_total"],
            ),
            "warm_speedup": _ratio(
                double_timing["warm_median"],
                single_timing["warm_median"],
            ),
            "temporary_memory_ratio": _ratio(
                double_memory["temp_size_in_bytes"],
                single_memory["temp_size_in_bytes"],
            ),
            "temporary_bytes_saved": (
                _difference(
                    double_memory["temp_size_in_bytes"],
                    single_memory["temp_size_in_bytes"],
                )
            ),
            "peak_memory_ratio": _ratio(
                double_memory["peak_memory_in_bytes"],
                single_memory["peak_memory_in_bytes"],
            ),
            "peak_bytes_saved": (
                _difference(
                    double_memory["peak_memory_in_bytes"],
                    single_memory["peak_memory_in_bytes"],
                )
            ),
        },
    }

block_until_ready(tree)

Synchronize every array leaf in a PyTree.

Source code in jaxquantum/utils/benchmarking.py

def block_until_ready(tree) -> None:
    """Synchronize every array leaf in a PyTree."""
    for leaf in jax.tree.leaves(tree):
        if hasattr(leaf, "block_until_ready"):
            leaf.block_until_ready()

jax_device_memory_stats()

Return allocator statistics reported by each JAX device.

Source code in jaxquantum/utils/benchmarking.py

def jax_device_memory_stats() -> dict[str, dict[str, int] | None]:
    """Return allocator statistics reported by each JAX device."""
    output = {}
    for device in jax.devices():
        stats = device.memory_stats()
        output[str(device)] = (
            None
            if stats is None
            else {
                key: int(value)
                for key, value in stats.items()
                if isinstance(value, (int, np.integer))
            }
        )
    return output

jax_hlo(function, *args, jit_kwargs=None, call_kwargs=None)

Return lowered StableHLO text for a function call.

Source code in jaxquantum/utils/benchmarking.py

def jax_hlo(
    function: Callable,
    *args,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
) -> str:
    """Return lowered StableHLO text for a function call."""
    return lower_jax_function(
        function,
        *args,
        jit_kwargs=jit_kwargs,
        call_kwargs=call_kwargs,
    ).as_text()

jax_memory_stats(compiled)

Return XLA's compiled buffer-size estimates.

Source code in jaxquantum/utils/benchmarking.py

def jax_memory_stats(compiled) -> dict[str, int | None]:
    """Return XLA's compiled buffer-size estimates."""
    memory = compiled.memory_analysis()
    if memory is None:
        return dict.fromkeys(_MEMORY_FIELDS)
    return {field: getattr(memory, field, None) for field in _MEMORY_FIELDS}

lower_jax_function(function, *args, jit_kwargs=None, call_kwargs=None)

Lower a function with the supplied JIT and call arguments.

Source code in jaxquantum/utils/benchmarking.py

def lower_jax_function(
    function: Callable,
    *args,
    jit_kwargs: Mapping[str, Any] | None = None,
    call_kwargs: Mapping[str, Any] | None = None,
):
    """Lower a function with the supplied JIT and call arguments."""
    return jax.jit(function, **dict(jit_kwargs or {})).lower(
        *args,
        **dict(call_kwargs or {}),
    )