# Peak FLOPs from parts: MAC cells x 2 ops per cell per cycle x clock. The
# two is one multiply and one add, and it is the only term here that is not
# looked up.
OPS_PER_MAC = 2


def peak_tflops(macs, clock_hz):
    return macs * OPS_PER_MAC * clock_hz / 1e12


# v1: one 256x256 systolic array.          ISCA 2017, section 2
# clock 700 MHz, published 92 TOPS.        ISCA 2017, Table 2
v1_macs = 256 * 256
reconcile("TPU v1 peak (int8)", peak_tflops(v1_macs, 700e6), 92.0, "TOPS")

# v4: 2 TensorCores x 4 MXUs, each 128x128. ISCA 2023, section 2
# clock 1050 MHz, published 275 TFLOPS.     ISCA 2023, Table 4
v4_macs = 2 * 4 * 128 * 128
v4_peak = peak_tflops(v4_macs, 1050e6)
reconcile("TPU v4 peak (bf16)", v4_peak, 275.0, "TFLOPS")

# One more zoom, same multiplication: a full v4 pod is 4096 chips, and
# Google advertises 1.1 exaflops for it. Two significant figures on the
# claim, so the tolerance has to be loose enough to read them.
# Google Cloud TPU v4 page, via /l/tpu/tensorcore-complex
reconcile("TPU v4 pod, 4096 chips", v4_peak * 4096 / 1e6, 1.1, "EFLOPS",
          tol=0.05, note="advertised to two figures")

print()
print(f"v1 and v4 hold the same cells twice over: {v1_macs} then {v4_macs}")

# v6e: 1 TensorCore x 2 MXUs (Google v6e page), each 256x256 (Cloud TPU
# architecture page). Published peak 918 TFLOPS bf16 (Google v6e page).
v6e_macs = 1 * 2 * 256 * 256
V6E_CLAIM = 918e12
implied_v6e = V6E_CLAIM / (v6e_macs * OPS_PER_MAC)

# yardsticks: the fastest clock Google itself has printed is v4 at 1050 MHz,
# and the fastest figure in the public record of any kind is 1.75 GHz for
# v5p (scaling book, secondary).
print(f"v6e cells from the published shape: {v6e_macs}")
print(f"implied clock: {implied_v6e / 1e9:.2f} GHz")
print(f"  against v4, printed:      {1050e6 / 1e9:.2f} GHz "
      f"({implied_v6e / 1050e6:.1f}x)")
print(f"  against v5p, secondary:   {1.75:.2f} GHz "
      f"({implied_v6e / 1.75e9:.1f}x)")
print()

# Two ways to make it close, neither of them evidence.
rescue_macs = 4 * 256 * 256          # assume four MXUs; Google's page says two
print(f"with 4 MXUs instead of 2:  {V6E_CLAIM / (rescue_macs * OPS_PER_MAC) / 1e9:.2f}"
      " GHz, which matches v5p exactly and contradicts the page")

# The same architecture page puts an MXU at 16K multiply-accumulate ops per
# cycle. 16K is 16,384, which is 128x128, not 256x256.
alt_macs = 2 * 16 * 1024
print(f"with the page's own 16K MACs per MXU: "
      f"{V6E_CLAIM / (alt_macs * OPS_PER_MAC) / 1e9:.1f} GHz, worse")
print(f"16K = {16 * 1024} = {int((16 * 1024) ** 0.5)}x{int((16 * 1024) ** 0.5)}, "
      "so the page disagrees with itself about the MXU shape")
print()

# Run the derivation the only honest way left: published shape, and the
# fastest clock anywhere in the public record. It comes up half short.
reconcile("TPU v6e peak at the fastest public clock",
          peak_tflops(v6e_macs, 1.75e9), 918.0, "TFLOPS",
          note="1.75 GHz is v5p's, and no v6e clock exists")

finding("tpu-v6e-open",
        f"The published v6e shape (2 MXUs at 256x256) and the published 918 "
        f"TFLOPS need {implied_v6e / 1e9:.2f} GHz, which no TPU approaches. "
        "Google publishes no v6e clock, and its architecture page gives an "
        "MXU MAC count that fits 128x128 rather than 256x256, so one input "
        "is wrong and nothing public says which.")

# A100 is the chip whose clock is printed. Whitepaper Table 3, p.39:
# 108 SMs, 1410 MHz boost, FP16 Tensor 312/624 (dense/sparse).
A100_SMS, A100_CLOCK, A100_FP16_DENSE = 108, 1.41e9, 312e12
a100_const = A100_FP16_DENSE / (A100_SMS * A100_CLOCK)
print(f"solve for the constant: {a100_const:.1f} FLOP/SM/clk -> call it 2048")
reconcile("A100 FP16 Tensor, dense", A100_SMS * 2048 * A100_CLOCK / 1e12,
          312.0, "TFLOPS")

# One sentence carries it to Hopper: "2x the MMA computational rates of the
# A100 SM on equivalent data types" (whitepaper p.22), per SM and per clock.
H100_PER_SM = 2048 * 2
H100_SMS = SM_COUNTS["H100 SXM5"]
h100_flops_per_clk = H100_SMS * H100_PER_SM
print(f"\nH100 SXM5: {H100_SMS} SMs x {H100_PER_SM} FLOP/SM/clk "
      f"= {h100_flops_per_clk:,} FLOP/clk")

clock_dense = BF16_DENSE * 1e12 / h100_flops_per_clk
clock_sparse = BF16_SPARSE * 1e12 / h100_flops_per_clk
print(f"clock implied by the dense rate:   {clock_dense / 1e9:.3f} GHz")
print(f"clock implied by the starred rate: {clock_sparse / 1e9:.3f} GHz  "
      "<- what the sparsity trap costs you")
print("third-party databases list 1830 MHz boost for H100 SXM5, which agrees;"
      "\nread that as a check on the arithmetic, not as the missing line")

# The datasheet rows that do not go through a Tensor Core, and the whitepaper
# preliminary column for the same part. FP32 is 128 lanes x 2 per SM, FP64 is
# 64 x 2, both from whitepaper Table 3.
fp32_per_clk = H100_SMS * FP32_PER_SM * 2
fp64_per_clk = H100_SMS * FP64_PER_SM * 2


def clock_band(printed, per_clk):
    "A printed rate is rounded, so it covers every clock in a band."
    lo = (printed - 0.5) * 1e12 / per_clk
    hi = (printed + 0.5) * 1e12 / per_clk
    return lo / 1e9, hi / 1e9


rows = [
    # rate on the page, TFLOPS, FLOPs per clock, source
    ("datasheet BF16 Tensor, dense", BF16_DENSE, h100_flops_per_clk, "nvidia.com H100"),
    ("datasheet FP32", 67.0, fp32_per_clk, "nvidia.com H100"),
    ("datasheet FP64", 34.0, fp64_per_clk, "nvidia.com H100"),
    ("preliminary FP32", 60.0, fp32_per_clk, "whitepaper Table 3"),
    ("preliminary FP64", 30.0, fp64_per_clk, "whitepaper Table 3"),
    ("preliminary FP16 Tensor, dense", 1000.0, h100_flops_per_clk, "whitepaper Table 3"),
]
for name, rate, per_clk, src in rows:
    lo, hi = clock_band(rate, per_clk)
    print(f"{name:<32} {rate:>7.1f} TF / {per_clk:>7,} = "
          f"{rate * 1e12 / per_clk / 1e9:.3f} GHz   band {lo:.3f}-{hi:.3f}  [{src}]")

# One row cannot be run through this arithmetic at all. FP64 Tensor Core
# prints 67 TFLOPS, the FP32 number over again, and NVIDIA never prints an
# FP64 rate per clock through the Tensor Cores.
print(f"\n{'datasheet FP64 Tensor Core':<32} {67.0:>7.1f} TF / "
      f"{'not published':>13} = cannot be derived")

# Do the two non-Tensor bands overlap, and does the Tensor clock sit in them?
lo32, hi32 = clock_band(67.0, fp32_per_clk)
lo64, hi64 = clock_band(34.0, fp64_per_clk)
overlap = (max(lo32, lo64), min(hi32, hi64))
print(f"\nFP32 and FP64 agree on a band: {overlap[0]:.3f}-{overlap[1]:.3f} GHz")
print(f"the Tensor rows demand {clock_dense / 1e9:.3f} GHz, which is "
      f"{'inside' if overlap[0] <= clock_dense / 1e9 <= overlap[1] else 'outside'}"
      " that band")
print("printed clock, either document: none. The whitepaper says 'Not "
      "Finalized' and the datasheet page lists no clock at all.")

# Assume the whole page is quoted at one clock, take the Tensor rows' clock,
# and predict the FP32 row from it. The miss is the size of the problem.
print()
reconcile("H100 FP32 at the Tensor rows' clock",
          fp32_per_clk * clock_dense / 1e12, 67.0, "TFLOPS",
          note="one page, one clock, assumed")

finding("h100-three-clocks",
        f"One page implies three clocks. The Tensor rows give "
        f"{clock_dense / 1e9:.3f} GHz, FP32 and FP64 agree on a band of "
        f"{overlap[0]:.3f} to {overlap[1]:.3f} GHz once rounding is allowed, "
        "and NVIDIA prints no clock in either document. A peak-FLOPS table "
        "is not necessarily quoted at one clock.")
finding("h100-fp64-tensor",
        "The FP64 Tensor Core row prints 67 TFLOPS and no FLOPs-per-clock "
        "figure exists for it, so its divisor can only be inferred from the "
        "doubling against plain FP64. Recorded as unavailable, not filled in.")