import torch
from torch import Tensor
import torch.utils.benchmark as benchmark


MIN_RUN_TIME_S = 5.0


@torch.library.custom_op("benchlib::clone_custom_op", mutates_args=())
def clone_custom_op(
    a: Tensor, b: Tensor, c: Tensor, d: Tensor, e: Tensor,
    f: Tensor, g: Tensor, h: Tensor, i: Tensor, j: Tensor,
    k: Tensor, l: Tensor, m: Tensor, n: Tensor,
) -> Tensor:
    return a.clone()


@torch.library.custom_op("benchlib::mutate_op", mutates_args=("a",))
def mutate_op(
    a: Tensor, b: Tensor, c: Tensor, d: Tensor, e: Tensor,
    f: Tensor, g: Tensor, h: Tensor, i: Tensor, j: Tensor,
    k: Tensor, l: Tensor, m: Tensor, n: Tensor,
) -> Tensor:
    a.zero_()
    return a.clone()


lib = torch.library.Library("benchlib2", "DEF")
lib.define(
    "clone_define_impl(Tensor a, Tensor b, Tensor c, Tensor d, Tensor e,"
    " Tensor f, Tensor g, Tensor h, Tensor i, Tensor j,"
    " Tensor k, Tensor l, Tensor m, Tensor n) -> Tensor"
)
lib.impl("clone_define_impl", lambda a, *_: a.clone(), "CPU")
lib.impl("clone_define_impl", lambda a, *_: a.clone(), "CUDA")


def clone_direct(a, b, c, d, e, f, g, h, i, j, k, l, m, n):
    return a.clone()


def bench(fn, args):
    t = benchmark.Timer(stmt="fn(*args)", globals={"fn": fn, "args": args})
    m = t.blocked_autorange(min_run_time=MIN_RUN_TIME_S)
    return m.median * 1e6


def bench_inference_mode(fn):
    t = benchmark.Timer(
        stmt="fn(*args)",
        globals={"fn": fn},
        setup=(
            "ctx = torch.inference_mode(); ctx.__enter__(); "
            "args = [torch.randn(4) for _ in range(14)]"
        ),
    )
    m = t.blocked_autorange(min_run_time=MIN_RUN_TIME_S)
    return m.median * 1e6


def main():
    tensors = [torch.randn(4) for _ in range(14)]
    grad_tensors = [torch.randn(4, requires_grad=True) for _ in range(14)]
    op_define = torch.ops.benchlib2.clone_define_impl
    op_custom_via_ops = torch.ops.benchlib.clone_custom_op

    print(f"PyTorch {torch.__version__}")
    print(f"{'Scenario':30s} {'us':>8s}")
    print("-" * 40)
    for label, fn, args in [
        ("custom_op (no grad)", clone_custom_op, tensors),
        ("custom_op via ops.*", op_custom_via_ops, tensors),
        ("custom_op (grad)", clone_custom_op, grad_tensors),
        ("mutate_op (no grad)", mutate_op, tensors),
        ("define/impl", op_define, tensors),
        ("direct call", clone_direct, tensors),
    ]:
        us = bench(fn, args)
        print(f"{label:30s} {us:8.1f}")

    us = bench_inference_mode(clone_custom_op)
    print(f"{'custom_op (inference_mode)':30s} {us:8.1f}")


if __name__ == "__main__":
    main()