# Benchmark asynchronous work correctly
Phase 08: Performance diagnosis · about 80 minutes · CPU
## What you will be able to do
- Draw the interval a timing actually measures.
- Exclude initial placement from a device-resident benchmark.
- Check an independent NumPy result before comparing timings.
- Produce a report containing first-call and warmed sample distributions.
## The problem
A first timing suggests a tenfold speedup. Before celebrating, let’s ask when the timer started and stopped. We’ll separate the first call from warm calls, wait for JAX to finish, and compare numerical results. By the end, you’ll have a timing report another learner can reproduce.
## The idea
JAX can return an array handle before device work finishes. A benchmark must define what starts the interval and what counts as completion. First-call preparation, warm execution and end-to-end requests answer different performance questions.
## Time a completed result
If a timer stops as soon as dispatch returns, it may exclude most of the computation. Blocking on the result inside the interval includes completion of the work required to produce that result.
Warm up the intended specialization before measuring repeated execution, and retain startup separately if it matters to the application. Include preprocessing and transfers in a separately named end-to-end measurement when they belong to the request.
The first-call and warm bars represent different intervals. Their ratio depends on the workload and environment. A short dispatch time is not evidence of device throughput, and a warm kernel timing does not describe an entire service request.
$$
T_{\text{warm}} = \operatorname{median}_{r=1 \dots R}\left(t_{\text{ready}}^{(r)} - t_{\text{start}}^{(r)}\right), \qquad T_{\text{first}} = T_{\text{compile}} + T_{\text{warm}}
$$
### Time the completed result
**Predict:** Where should synchronization occur in completed-result latency measurement?
*Architecture and dataflow mechanism diagram.*
The host dispatches and waits while device work proceeds. The measurement lane includes readiness. Lengths illustrate ordering only and are not measured durations. Compilation and data preparation need separately chosen boundaries when included in a benchmark.
### Pause and reason
Where should synchronization occur in completed-result latency measurement?
Compare your reasoning
Before stopping the timer. Synchronizing afterward can omit unfinished work from the reported interval.
Reference solution
```python
# Exercise solution: Repeat the benchmark with 96 rows while preserving the feature and...
larger_host = rng.normal(size=(96, 32)).astype(np.float32)
# Place `larger` explicitly onto the target JAX device.
larger = jax.device_put(larger_host, cpu)
# Synchronize host execution until asynchronous device computation completes.
larger.block_until_ready()
# Synchronize host execution until asynchronous device computation completes.
compiled_predict(larger, w).block_until_ready()
larger_result, larger_times = synchronized_samples(compiled_predict, (larger, w), 5)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(larger_result), np.tanh(larger_host @ w_host), atol=2e-5, rtol=2e-5)
print("Changed workload:", larger.shape, larger_times)
```
Hint
Square before reducing the output-feature axis; the result has shape $(64,)$.
Reference solution and reasoning
```python
# Benchmark a different equivalent workload (Transfer): Changing the computation changes the performance question.
# Wrap with `jax.jit` (`energy`) so XLA traces and compiles the function.
energy = jax.jit(lambda a, b: jnp.sum((a @ b)**2, axis=1))
# Synchronize host execution until asynchronous device computation completes.
energy(x, w).block_until_ready()
energy_result, energy_times = synchronized_samples(energy, (x, w), 5)
# Reduce along axis=1 to compute `energy_reference`.
energy_reference = np.sum((x_host @ w_host)**2, axis=1)
# Check tensor shape invariant: `energy_result.shape == (64,)`
assert energy_result.shape == (64,)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(energy_result), energy_reference, atol=2e-3, rtol=2e-5)
print("Energy shape and samples:", energy_result.shape, energy_times)
```
Changing the computation changes the performance question. The shape and independent numerical check prevent a benchmark from accidentally timing an unrelated scalar reduction.
Hint
Prepare and wait for inputs, reuse one compiled callable, warm it, then synchronize the returned output inside every measured interval.
Reference solution and reasoning
```python
# Repair a timer that includes setup but omits completion (Diagnosis): The repaired boundary excludes setup, separates warmup and...
def repaired_benchmark(fn, args):
# Synchronize host execution until asynchronous device computation completes.
jax.block_until_ready(args)
# Synchronize host execution until asynchronous device computation completes.
jax.block_until_ready(fn(*args))
result, measured = synchronized_samples(fn, args, 5)
return {"boundary": "placed inputs; warm call plus output wait", "samples": measured}, result
fixed_report, fixed_result = repaired_benchmark(compiled_predict, (x, w))
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(fixed_result), reference, atol=2e-5, rtol=2e-5)
assert len(fixed_report["samples"]) == 5
print(fixed_report)
```
The repaired boundary excludes setup, separates warmup and waits before stopping. A speedup cannot be reconstructed from the flawed interval; rerun both implementations under the same contract.
Answer and explanation
A warmed call including dispatch and waiting for output readiness, excluding prior input placement
The boundary includes the call and its completion wait. It excludes input placement and does not isolate compiler or kernel stages.