Coverage for jaxquantum/utils/benchmarking.py: 91%
149 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-27 22:28 +0000
1"""Synchronized JAX compilation, execution, and memory measurements."""
3from __future__ import annotations
5import platform
6import statistics
7import time
8from collections.abc import Callable, Mapping
9from typing import Any
11import jax
12import numpy as np
14_MEMORY_FIELDS = (
15 "argument_size_in_bytes",
16 "output_size_in_bytes",
17 "alias_size_in_bytes",
18 "temp_size_in_bytes",
19 "peak_memory_in_bytes",
20 "generated_code_size_in_bytes",
21 "host_argument_size_in_bytes",
22 "host_output_size_in_bytes",
23 "host_alias_size_in_bytes",
24 "host_temp_size_in_bytes",
25)
28def block_until_ready(tree) -> None:
29 """Synchronize every array leaf in a PyTree."""
30 for leaf in jax.tree.leaves(tree):
31 if hasattr(leaf, "block_until_ready"):
32 leaf.block_until_ready()
35def lower_jax_function(
36 function: Callable,
37 *args,
38 jit_kwargs: Mapping[str, Any] | None = None,
39 call_kwargs: Mapping[str, Any] | None = None,
40):
41 """Lower a function with the supplied JIT and call arguments."""
42 return jax.jit(function, **dict(jit_kwargs or {})).lower(
43 *args,
44 **dict(call_kwargs or {}),
45 )
48def jax_memory_stats(compiled) -> dict[str, int | None]:
49 """Return XLA's compiled buffer-size estimates."""
50 memory = compiled.memory_analysis()
51 if memory is None:
52 return dict.fromkeys(_MEMORY_FIELDS)
53 return {field: getattr(memory, field, None) for field in _MEMORY_FIELDS}
56def jax_device_memory_stats() -> dict[str, dict[str, int] | None]:
57 """Return allocator statistics reported by each JAX device."""
58 output = {}
59 for device in jax.devices():
60 stats = device.memory_stats()
61 output[str(device)] = (
62 None
63 if stats is None
64 else {
65 key: int(value)
66 for key, value in stats.items()
67 if isinstance(value, (int, np.integer))
68 }
69 )
70 return output
73def jax_hlo(
74 function: Callable,
75 *args,
76 jit_kwargs: Mapping[str, Any] | None = None,
77 call_kwargs: Mapping[str, Any] | None = None,
78) -> str:
79 """Return lowered StableHLO text for a function call."""
80 return lower_jax_function(
81 function,
82 *args,
83 jit_kwargs=jit_kwargs,
84 call_kwargs=call_kwargs,
85 ).as_text()
88def _percentiles(samples):
89 if len(samples) < 2:
90 return samples[0], samples[0]
91 deciles = statistics.quantiles(samples, n=10, method="inclusive")
92 return deciles[0], deciles[-1]
95def _scalar_costs(costs):
96 if not costs:
97 return {}
98 output = {}
99 for key, value in costs.items():
100 try:
101 output[key] = value.item()
102 except AttributeError:
103 output[key] = value
104 return output
107def _benchmark_once(
108 function: Callable,
109 args,
110 iterations,
111 warmup,
112 clear_caches,
113 include_hlo,
114 jit_kwargs,
115 call_kwargs,
116):
117 if clear_caches:
118 jax.clear_caches()
119 device_memory_before = jax_device_memory_stats()
121 start = time.perf_counter()
122 lowered = lower_jax_function(
123 function,
124 *args,
125 jit_kwargs=jit_kwargs,
126 call_kwargs=call_kwargs,
127 )
128 lowering_s = time.perf_counter() - start
130 start = time.perf_counter()
131 compiled = lowered.compile()
132 compilation_s = time.perf_counter() - start
134 start = time.perf_counter()
135 first_output = compiled(*args, **call_kwargs)
136 block_until_ready(first_output)
137 first_execution_s = time.perf_counter() - start
139 for _ in range(warmup):
140 block_until_ready(compiled(*args, **call_kwargs))
142 samples = []
143 for _ in range(iterations):
144 start = time.perf_counter()
145 block_until_ready(compiled(*args, **call_kwargs))
146 samples.append(time.perf_counter() - start)
148 p10, p90 = _percentiles(samples)
149 hlo = lowered.as_text()
150 report = {
151 "function": getattr(function, "__qualname__", repr(function)),
152 "platform": platform.platform(),
153 "jax": jax.__version__,
154 "backend": jax.default_backend(),
155 "devices": [str(device) for device in jax.devices()],
156 "timings_s": {
157 "lowering": lowering_s,
158 "compilation": compilation_s,
159 "first_execution": first_execution_s,
160 "cold_total": lowering_s + compilation_s + first_execution_s,
161 "warm_median": statistics.median(samples),
162 "warm_min": min(samples),
163 "warm_max": max(samples),
164 "warm_p10": p10,
165 "warm_p90": p90,
166 },
167 "memory_bytes": jax_memory_stats(compiled),
168 "device_memory_stats": {
169 "before": device_memory_before,
170 "after": jax_device_memory_stats(),
171 },
172 "cost_analysis": _scalar_costs(compiled.cost_analysis()),
173 "hlo": {
174 "characters": len(hlo),
175 "lines": hlo.count("\n") + 1,
176 },
177 "iterations": iterations,
178 "warmup": warmup,
179 }
180 if include_hlo:
181 report["hlo"]["text"] = hlo
182 return report, first_output
185def _cast_precision(tree, dtype, complex_dtype):
186 def cast(leaf):
187 leaf_dtype = getattr(leaf, "dtype", None)
188 if leaf_dtype is None:
189 return leaf
190 try:
191 if np.issubdtype(leaf_dtype, np.complexfloating):
192 return leaf.astype(complex_dtype)
193 if np.issubdtype(leaf_dtype, np.floating):
194 return leaf.astype(dtype)
195 except TypeError:
196 pass
197 return leaf
199 return jax.tree.map(cast, tree)
202def _accuracy_stats(reference, candidate):
203 if jax.tree.structure(reference) != jax.tree.structure(candidate):
204 raise ValueError("precision outputs have different PyTree structures")
205 reference_leaves = jax.tree.leaves(reference)
206 candidate_leaves = jax.tree.leaves(candidate)
208 max_absolute = 0.0
209 max_relative = 0.0
210 squared_error = 0.0
211 squared_reference = 0.0
212 elements = 0
213 for reference_leaf, candidate_leaf in zip(
214 reference_leaves,
215 candidate_leaves,
216 ):
217 reference_array = np.asarray(jax.device_get(reference_leaf))
218 candidate_array = np.asarray(jax.device_get(candidate_leaf))
219 if reference_array.shape != candidate_array.shape:
220 raise ValueError("precision outputs have different shapes")
221 if not np.issubdtype(reference_array.dtype, np.inexact):
222 continue
223 difference = np.abs(reference_array - candidate_array)
224 magnitude = np.abs(reference_array)
225 if difference.size:
226 max_absolute = max(max_absolute, float(np.max(difference)))
227 threshold = np.finfo(reference_array.real.dtype).tiny
228 relative = difference / np.maximum(magnitude, threshold)
229 max_relative = max(max_relative, float(np.max(relative)))
230 squared_error += float(np.sum(difference**2))
231 squared_reference += float(np.sum(magnitude**2))
232 elements += difference.size
234 relative_l2 = (
235 np.sqrt(squared_error / squared_reference)
236 if squared_reference
237 else np.sqrt(squared_error)
238 )
239 return {
240 "max_absolute_error": max_absolute,
241 "max_relative_error": max_relative,
242 "relative_l2_error": float(relative_l2),
243 "elements_compared": elements,
244 }
247def _ratio(numerator, denominator):
248 if numerator is None or denominator in (None, 0):
249 return None
250 return numerator / denominator
253def _difference(left, right):
254 if left is None or right is None:
255 return None
256 return left - right
259def benchmark_precision(
260 function: Callable,
261 *args,
262 iterations: int = 25,
263 warmup: int = 1,
264 clear_caches: bool = True,
265 include_hlo: bool = False,
266 jit_kwargs: Mapping[str, Any] | None = None,
267 call_kwargs: Mapping[str, Any] | None = None,
268) -> dict[str, Any]:
269 """Compare double and single precision performance and output accuracy."""
270 if iterations < 1 or warmup < 0:
271 raise ValueError("iterations must be positive and warmup non-negative")
272 original_x64 = jax.config.x64_enabled
273 reports = {}
274 outputs = {}
275 try:
276 for name, enabled, real_dtype, complex_dtype in (
277 ("double", True, np.float64, np.complex128),
278 ("single", False, np.float32, np.complex64),
279 ):
280 jax.config.update("jax_enable_x64", enabled)
281 precision_args = _cast_precision(args, real_dtype, complex_dtype)
282 precision_kwargs = _cast_precision(
283 dict(call_kwargs or {}),
284 real_dtype,
285 complex_dtype,
286 )
287 reports[name], output = _benchmark_once(
288 function,
289 precision_args,
290 iterations,
291 warmup,
292 clear_caches,
293 include_hlo,
294 jit_kwargs,
295 precision_kwargs,
296 )
297 outputs[name] = jax.device_get(output)
298 del output
299 finally:
300 jax.config.update("jax_enable_x64", original_x64)
301 if clear_caches:
302 jax.clear_caches()
304 double_timing = reports["double"]["timings_s"]
305 single_timing = reports["single"]["timings_s"]
306 double_memory = reports["double"]["memory_bytes"]
307 single_memory = reports["single"]["memory_bytes"]
308 return {
309 "double": reports["double"],
310 "single": reports["single"],
311 "accuracy": _accuracy_stats(outputs["double"], outputs["single"]),
312 "single_vs_double": {
313 "cold_speedup": _ratio(
314 double_timing["cold_total"],
315 single_timing["cold_total"],
316 ),
317 "warm_speedup": _ratio(
318 double_timing["warm_median"],
319 single_timing["warm_median"],
320 ),
321 "temporary_memory_ratio": _ratio(
322 double_memory["temp_size_in_bytes"],
323 single_memory["temp_size_in_bytes"],
324 ),
325 "temporary_bytes_saved": (
326 _difference(
327 double_memory["temp_size_in_bytes"],
328 single_memory["temp_size_in_bytes"],
329 )
330 ),
331 "peak_memory_ratio": _ratio(
332 double_memory["peak_memory_in_bytes"],
333 single_memory["peak_memory_in_bytes"],
334 ),
335 "peak_bytes_saved": (
336 _difference(
337 double_memory["peak_memory_in_bytes"],
338 single_memory["peak_memory_in_bytes"],
339 )
340 ),
341 },
342 }
345def benchmark_jax_function(
346 function: Callable,
347 *args,
348 iterations: int = 25,
349 warmup: int = 1,
350 clear_caches: bool = True,
351 include_hlo: bool = False,
352 compare_precision: bool = False,
353 jit_kwargs: Mapping[str, Any] | None = None,
354 call_kwargs: Mapping[str, Any] | None = None,
355) -> dict[str, Any]:
356 """Collect synchronized JAX timing, memory, HLO, cost, and precision stats."""
357 if compare_precision:
358 return benchmark_precision(
359 function,
360 *args,
361 iterations=iterations,
362 warmup=warmup,
363 clear_caches=clear_caches,
364 include_hlo=include_hlo,
365 jit_kwargs=jit_kwargs,
366 call_kwargs=call_kwargs,
367 )
368 if iterations < 1 or warmup < 0:
369 raise ValueError("iterations must be positive and warmup non-negative")
370 return _benchmark_once(
371 function,
372 args,
373 iterations,
374 warmup,
375 clear_caches,
376 include_hlo,
377 jit_kwargs,
378 dict(call_kwargs or {}),
379 )[0]