Four numbers for one call

The chapter builds a harness by hand, and it is the right thing to build once, because writing it is how you learn what it defends against. Run four variants of it against one nn.Linear(512, 512) on a batch of 64 and the first column disagrees with the other three by a factor of seven.

A single cold call measured 1.451 milliseconds on the first of five runs, and stayed between 1.015 and 1.740 across all five. The warm loop in the same script never rose above 0.234. So a first call costs four to eight times the steady figure here, and none of that difference is compute: it is one-time allocation, the first pass through a fresh dispatch path, and whatever the operating system does the first time a page gets touched.

The other three columns are the interesting part, because they agree. The cold loop, the warm loop and torch.utils.benchmark.Timer landed inside 0.194 to 0.250 milliseconds on every one of the five runs, and the run-to-run spread of each one is about 20 percent. So the choice between a hand loop and the shipped instrument is not a question of accuracy on this program. It is a question of which one tells you the spread instead of hiding it, which is the next section.

A cold call and a warm loop are two different measurements of one function.

import time

import torch
import torch.nn as nn
import torch.utils.benchmark as benchmark

torch.manual_seed(0)
lin = nn.Linear(512, 512)
x = torch.randn(64, 512)

def one_call():
    t0 = time.perf_counter()
    lin(x)
    return (time.perf_counter() - t0) * 1e3

def loop(n=50):
    t0 = time.perf_counter()
    for _ in range(n):
        lin(x)
    return (time.perf_counter() - t0) / n * 1e3

first = one_call()
cold_loop = loop()
for _ in range(50):
    lin(x)
warm_loop = loop()
m = benchmark.Timer(stmt="lin(x)", globals={"lin": lin, "x": x},
                    num_threads=torch.get_num_threads()).blocked_autorange(min_run_time=2.0)

print(f"one cold call         {first:7.3f} ms")
print(f"50-call mean, cold    {cold_loop:7.3f} ms")
print(f"50-call mean, warm    {warm_loop:7.3f} ms")
print(f"Timer median          {m.median * 1e3:7.3f} ms   iqr {m.iqr * 1e3:6.3f}")

# ---- stdout ----
# one cold call           1.451 ms
# 50-call mean, cold      0.207 ms
# 50-call mean, warm      0.196 ms
# Timer median            0.197 ms   iqr  0.049

torch.utils.benchmark.Timer is not a nicer wrapper around perf_counter. It picks a block size by warming up until the timer call itself accounts for under 0.1 percent of the measurement, runs blocks until it has spent min_run_time, and returns a Measurement that carries every block time rather than one average. That is what a hand loop cannot give you back, because a hand loop has thrown the distribution away by the time it divides.

Point it at one training step of a small MLP here and it answers 1.39 milliseconds, then says on the next line that the interquartile range is 41.6 percent of that. The warning is generated, not written by a human: common.py sets two thresholds, a tenth and a quarter of the median, and picks between two sentences depending on which one the measurement crossed. Four runs at eight threads crossed the higher one every time.

The same object reports significant_figures as 2, estimated from the interquartile region alone so the tails cannot flatter it. Two digits is what this measurement is entitled to. Quoting 1.394 milliseconds off this run means publishing noise in the third digit.

Provenance ages, which is the other reason to report the spread. The chapter's harness measured 0.142 milliseconds for a Linear of this shape on the day it was written; the same shape measures 0.194 to 0.250 today on the same machine. Both are honest. A wall-clock figure is a fact about a machine on a day, so a comparison belongs inside one run and never across two.

$ cat step_timer.py
import torch
import torch.nn as nn
import torch.utils.benchmark as benchmark

torch.manual_seed(0)
model = nn.Sequential(nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 10))
opt = torch.optim.SGD(model.parameters(), lr=0.1)
x, y = torch.randn(64, 512), torch.randint(0, 10, (64,))
loss_fn = nn.CrossEntropyLoss()

def step():
    opt.zero_grad(set_to_none=True)
    loss_fn(model(x), y).backward()
    opt.step()

m = benchmark.Timer(stmt="step()", globals={"step": step},
                    label="one training step", sub_label="64 x 512 MLP",
                    num_threads=torch.get_num_threads()).blocked_autorange(min_run_time=3.0)
print(m)
print("has_warnings:", m.has_warnings, "| significant_figures:", m.significant_figures)

$ python3 step_timer.py
one training step: 64 x 512 MLP
  Median: 1.39 ms
  IQR:    0.58 ms (1.18 to 1.76)
  1898 measurements, 1 runs per measurement, 8 threads
  WARNING: Interquartile range is 41.6% of the median measurement.
           This suggests significant environmental influence.
has_warnings: True | significant_figures: 2

Timer.__init__ takes num_threads: int = 1, and that default is the most common way a torch benchmark ends up describing a program nobody runs. This process has torch.get_num_threads() at 8. Construct a Timer without saying otherwise and it measures your matmul with seven of those threads idle.

The gap is not small. A 512 by 512 matmul measured 4072 microseconds at one thread and 1122 at eight, and the default-constructed Timer in the same script reported 4089, sitting with the one-thread column exactly as the signature says it should. Read that as the cost inside an eight-thread training loop and you are off by a factor of 3.6.

Threading also moves the spread, in the direction you would not guess. The single-threaded runs of the training step above had the tighter interquartile range, 22.7 and 24.3 percent against 34.6 to 42.3 at eight threads, and that is the difference between the two warning sentences the last section showed. Eight threads on a laptop contend with whatever else the laptop is doing. Fewer threads measured slower and measured steadier.

So the thread count is part of the measurement, in the same way the shapes and the dtype are. Say it next to the number. torch.set_num_threads sets the intra-op pool for the process; num_threads on a Timer sets it for that measurement only and restores it after.

import torch
import torch.utils.benchmark as benchmark

print("torch.get_num_threads()", torch.get_num_threads())
x = torch.randn(512, 512)
default = benchmark.Timer(stmt="x @ x", globals={"x": x})
print("Timer default num_threads:", default._task_spec.num_threads)
r = default.blocked_autorange(min_run_time=1.0)
print(f"default timer  median {r.median * 1e6:8.1f} us")
for nt in (1, 2, 4, 8):
    m = benchmark.Timer(stmt="x @ x", globals={"x": x}, num_threads=nt).blocked_autorange(min_run_time=1.0)
    print(f"threads={nt}      median {m.median * 1e6:8.1f} us  iqr {m.iqr * 1e6:7.1f}")

# ---- stdout ----
# torch.get_num_threads() 8
# Timer default num_threads: 1
# default timer  median   4088.8 us
# threads=1      median   4071.7 us  iqr   459.5
# threads=2      median   2297.9 us  iqr   430.9
# threads=4      median   1331.9 us  iqr   123.7
# threads=8      median   1122.0 us  iqr    97.7

import torch
import torch.nn as nn
import torch.utils.benchmark as benchmark

torch.manual_seed(0)
lin = nn.Linear(512, 512)
results = []
for batch in (1, 16, 64, 256):
    x = torch.randn(batch, 512)
    for nt in (1, 8):
        results.append(
            benchmark.Timer(
                stmt="lin(x)",
                globals={"lin": lin, "x": x},
                label="nn.Linear(512, 512)",
                sub_label=f"batch {batch}",
                description=f"{nt} thread" + ("s" if nt > 1 else ""),
                num_threads=nt,
            ).blocked_autorange(min_run_time=1.0)
        )
benchmark.Compare(results).print()

# ---- stdout ----
# [--------- nn.Linear(512, 512) ----------]
#                  |  1 thread  |  8 threads
# 1 threads: -------------------------------
#       batch 1    |     52.6   |
#       batch 16   |    198.5   |
#       batch 64   |    594.8   |
#       batch 256  |   2252.0   |
# 8 threads: -------------------------------
#       batch 1    |            |     42.1
#       batch 16   |            |     88.5
#       batch 64   |            |    198.8
#       batch 256  |            |    687.3
#
# Times are in microseconds (us).

Underneath all of this there is a cost that does not shrink when the work does. One elementwise add on a one-element tensor took 8.70 microseconds here. On 64 elements, 8.94. On 4096 elements, 9.84. The tensor grew by four thousand times and the call got 13 percent more expensive.

The Python around it is not the cost. An empty function call in the same loop measured 0.068 microseconds and x.dim() measured 0.112, so what the 9 microseconds buys is the dispatch itself: the operator lookup, the shape and dtype checks, the output allocation, the kernel entry. In-place add_ came in at 6.96, and the missing two microseconds are roughly the allocation the out-of-place version pays.

Past 4096 elements the arithmetic finally takes over, 93.72 microseconds at 262144 and 729.96 at a million. The figure to carry out of this is the crossover, not the constant. Below a few thousand elements per op you are timing PyTorch's dispatcher, and no faster kernel can help you there.

That is the same wall the compile chapters attack from the other side. A model that spends its time in thousands of small ops is what torch.compile fuses, and the lesson arc at /pytorch/dynamo counts how many graphs you actually got. A benchmark that cannot see the floor cannot tell you whether fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → is worth attempting.

import time

import torch

torch.set_num_threads(1)
x1 = torch.ones(1)

def loop(fn, n=20000):
    fn()
    t0 = time.perf_counter()
    for _ in range(n):
        fn()
    return (time.perf_counter() - t0) / n * 1e6

def nothing():
    pass

print(f"python call only  {min(loop(nothing) for _ in range(5)):7.3f} us")
print(f"x.dim()           {min(loop(x1.dim) for _ in range(5)):7.3f} us")
print(f"x + 1.0           {min(loop(lambda: x1 + 1.0) for _ in range(5)):7.3f} us")
print(f"x.add_(0.0)       {min(loop(lambda: x1.add_(0.0)) for _ in range(5)):7.3f} us")

def sized(numel, n):
    x = torch.ones(numel)
    y = x + 1.0
    t0 = time.perf_counter()
    for _ in range(n):
        y = x + 1.0
    return (time.perf_counter() - t0) / n * 1e6

for numel in (1, 64, 4096, 262144, 1048576):
    n = 5000 if numel < 262144 else 300
    print(f"add on {numel:9d} elements {min(sized(numel, n) for _ in range(5)):9.2f} us")

# ---- stdout ----
# python call only    0.068 us
# x.dim()             0.112 us
# x + 1.0             9.120 us
# x.add_(0.0)         6.960 us
# add on         1 elements      8.70 us
# add on        64 elements      8.94 us
# add on      4096 elements      9.84 us
# add on    262144 elements     93.72 us
# add on   1048576 elements    729.96 us

Nothing measured above needed a synchronize, because CPU dispatch runs the op before the line returns. The chapter names the two backends where that stops being true, and the jax path's lesson at /jax/performance/what-the-clock-caught measures the async version of the same mistake in detail: what an unblocked loop actually catches, why blocking every call and blocking once report different true numbers, and when to quote a minimum against a mean. Those traps are told once, there, and they transfer to torch without modification.

What torch adds is that the instrument already handles the first of them. When the module is imported into a process where CUDA is built and available, timer is redefined to synchronize before reading the clock, and that function is the default timer argument of every Timer you construct. The harness the chapter writes by hand is the harness Timer already is, on whichever device you run it.

The lazy backend is the one case where no instrument saves you. A torch_xla tensor records instead of computing, so a timer around it measures graph construction until something forces a sync, and LAB·P5 is where that boundary gets counted call by call on real hardware.

if torch.backends.cuda.is_built() and torch.cuda.is_available():
    def timer() -> float:
        torch.cuda.synchronize()
        return timeit.default_timer()

# and from the Timer docstring, on the timer argument:
#
#     Callable which returns the current time. If PyTorch was built
#     without CUDA or there is no GPU present, this defaults to
#     `timeit.default_timer`; otherwise it will synchronize CUDA before
#     measuring the time.