"""Measure memory bandwidth on the accelerator and on the CPU, two ways, with PyTorch.

Purpose: turn the vendor's published bandwidth figure into a measurement you made
         yourself, with two workloads dominated by moving bytes rather than arithmetic.
Platform: all (cuda on Tracks S and N and on Track X with the ROCm build, mps on
          Track M, cpu everywhere as a comparison and as the fallback)
Minimum memory: 8 GB
Assumes: torch is installed in the active environment; each device has room for two
         buffers of --size-mb (1 GiB at the default), because a copy needs a source
         and a destination

Usage: python3 bandwidth-test.py [--size-mb 512] [--iters 20] [--repeats 3]
                                 [--device auto] [--labbook labbook.md]

Method. Two tests per device, each run --repeats times after a warm-up. The median of
the repeats is reported with the spread, (slowest - fastest) / median. 1 GB = 1e9 bytes.
  copy  dst.copy_(src) on a buffer of N bytes reads N bytes and writes N bytes, so it
        reports (2 x N x iters) / seconds: bus traffic in both directions.
  read  y = W @ x, with W a square float32 matrix of about N bytes and x a vector, is one
        decode step at batch 1 in miniature: it reads every weight once and writes only
        the small output vector, so it reports (bytes of W x iters) / seconds.
The read figure counts bytes the same way predict-decode.py counts bytes per token, so it
is the figure to pass to that script.
"""
import argparse
import json
import statistics
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

import torch

GB = 1_000_000_000
MIB = 1024 * 1024
SPREAD_WARNING = 10.0  # per cent


def pick_device(requested):
    if requested != "auto":
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def synchronize(device):
    """Block until the device has finished, so that timings are not of queueing."""
    if device.type == "cuda":
        torch.cuda.synchronize()
    elif device.type == "mps":
        torch.mps.synchronize()


def release(device):
    if device.type == "cuda":
        torch.cuda.empty_cache()
    elif device.type == "mps":
        torch.mps.empty_cache()


def describe(device):
    if device.type == "cuda":
        if torch.version.hip:
            runtime = f"ROCm (HIP {torch.version.hip})"
        else:
            runtime = f"CUDA {torch.version.cuda}"
        return f"{torch.cuda.get_device_name(0)}, {runtime}"
    if device.type == "mps":
        return "Apple silicon GPU through MPS"
    return f"CPU, {torch.get_num_threads()} threads"


def timed(device, fn, iters):
    synchronize(device)
    start = time.perf_counter()
    for _ in range(iters):
        fn()
    synchronize(device)
    return time.perf_counter() - start


def copy_test(device, size_mb, iters, repeats):
    nbytes = size_mb * MIB
    src = torch.empty(nbytes // 4, dtype=torch.float32, device=device)
    src.uniform_(-1.0, 1.0)
    dst = torch.empty_like(src)
    for _ in range(3):  # the first copies pay for allocation and kernel setup
        dst.copy_(src)
    rates = []
    for _ in range(repeats):
        seconds = timed(device, lambda: dst.copy_(src), iters)
        rates.append(2 * nbytes * iters / seconds / GB)
    del src, dst
    release(device)
    return rates, 2 * nbytes


def read_test(device, size_mb, iters, repeats):
    side = int((size_mb * MIB / 4) ** 0.5)
    w = torch.empty(side, side, dtype=torch.float32, device=device)
    w.uniform_(-1.0, 1.0)
    x = torch.empty(side, dtype=torch.float32, device=device)
    x.uniform_(-1.0, 1.0)
    nbytes = w.numel() * w.element_size()
    for _ in range(3):
        torch.matmul(w, x)
    rates = []
    for _ in range(repeats):
        seconds = timed(device, lambda: torch.matmul(w, x), iters)
        rates.append(nbytes * iters / seconds / GB)
    del w, x
    release(device)
    return rates, side, nbytes


def summarise(rates):
    median = statistics.median(rates)
    spread = (max(rates) - min(rates)) / median * 100.0
    return median, spread


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--size-mb", type=int, default=512,
                        help="bytes per buffer in MiB; the copy test allocates two")
    parser.add_argument("--iters", type=int, default=20,
                        help="operations per timed repeat")
    parser.add_argument("--repeats", type=int, default=3,
                        help="timed repeats per test; the median is reported")
    parser.add_argument("--device", default="auto", help="auto, cuda, mps or cpu")
    parser.add_argument("--labbook", default=None,
                        help="append one JSON line per run to this file")
    args = parser.parse_args()
    if args.size_mb < 1 or args.iters < 1 or args.repeats < 1:
        sys.exit("bandwidth-test: --size-mb, --iters and --repeats must be at least 1")

    device = pick_device(args.device)
    devices = [device] if device.type == "cpu" else [device, torch.device("cpu")]
    print(f"torch {torch.__version__}")
    for d in devices:
        print(f"{d.type}: {describe(d)}")
    print(f"buffer {args.size_mb} MiB; {args.iters} operations x {args.repeats} repeats "
          f"per test; 1 GB = 1e9 bytes\n")
    header = f"{'device':<7}{'test':<6}{'bytes per operation':<32}{'GB/s':>8}{'spread':>9}"
    print(header)
    print("-" * len(header))

    results = {}
    for d in devices:
        try:
            copy_rates, copy_bytes = copy_test(d, args.size_mb, args.iters, args.repeats)
            read_rates, side, read_bytes = read_test(d, args.size_mb, args.iters, args.repeats)
        except RuntimeError as exc:
            reason = str(exc).splitlines()[0][:160]
            sys.exit(f"bandwidth-test: {d.type} failed with --size-mb {args.size_mb}: {reason}\n"
                     f"bandwidth-test: halve --size-mb and run again")
        copy_median, copy_spread = summarise(copy_rates)
        read_median, read_spread = summarise(read_rates)
        print(f"{d.type:<7}{'copy':<6}{f'{copy_bytes:,} read + written':<32}"
              f"{copy_median:>8.1f}{copy_spread:>8.1f}%")
        print(f"{d.type:<7}{'read':<6}{f'{read_bytes:,} read ({side}^2 x 4)':<32}"
              f"{read_median:>8.1f}{read_spread:>8.1f}%")
        results[d.type] = {
            "copy": round(copy_median, 1), "read": round(read_median, 1),
            "copy_runs": [round(r, 1) for r in copy_rates],
            "read_runs": [round(r, 1) for r in read_rates],
            "read_matrix_side": side,
        }
        if max(copy_spread, read_spread) > SPREAD_WARNING:
            print(f"        spread above {SPREAD_WARNING:.0f} per cent on {d.type}: "
                  f"the machine was not steady (see Troubleshooting)")

    if args.labbook:
        record = {
            "lab": "part-05/bandwidth-test",
            "date": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "tool": "torch", "torch": torch.__version__,
            "device": device.type, "device_name": describe(device),
            "method": "copy: bytes read + written per second; read: W @ x, bytes of W per second",
            "size_mb": args.size_mb, "iters": args.iters, "repeats": args.repeats,
            "gbps": results,
        }
        with Path(args.labbook).open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, sort_keys=True) + "\n")
        print(f"\nrecorded in {args.labbook}")


if __name__ == "__main__":
    main()