"""Proof: the fixed cost of one eager op, and why size hides it.

Times the same `a + b` at two sizes. The one-element add is nearly
pure machinery (dispatch, wrapping, allocation); the 4M-element add is
nearly pure arithmetic. CPU, single process.
"""
import time
import torch

def per_op_us(a, b, iters):
    # warmup
    for _ in range(2000):
        a + b
    t0 = time.perf_counter()
    for _ in range(iters):
        a + b
    return (time.perf_counter() - t0) / iters * 1e6

tiny = per_op_us(torch.ones(1), torch.ones(1), 200_000)
big_n = 4_000_000
big = per_op_us(torch.ones(big_n), torch.ones(big_n), 2_000)

print(f"torch {torch.__version__}, cpu")
print(f"add, 1 element   : {tiny:8.3f} us/op")
print(f"add, 4M elements : {big:8.3f} us/op")
print(f"machinery share of the tiny op: ~all of it")
print(f"ops/sec you can issue from python: {1e6/tiny:,.0f}")

# The sweep behind the toll meter: the same add at twelve sizes,
# 1 to 4M elements in powers of four. Every dot on the widget's
# axis is one line of this output.
import json
sweep = []
for k in range(12):
    n = 4 ** k
    iters = max(1_000, min(200_000, 40_000_000 // max(n, 1)))
    us = per_op_us(torch.ones(n), torch.ones(n), iters)
    sweep.append({"n": n, "us": round(us, 3)})
    print(f"add, {n:>9,} elements : {us:9.3f} us/op")
print("JSON_SWEEP=" + json.dumps(sweep))