import torch
import numpy as np
from tvm import ffi as tvm_ffi
import time


def print_speed(name, speed):
    print(f"{name:<40} {speed} sec/call")


def print_error(name, error):
    print(f"{name:<40} {error}")


def baseline_torch_add(repeat):
    """Run torch.add with one element"""

    def run_bench(device):
        x = torch.arange(1, device=device)
        y = torch.arange(1, device=device)
        z = torch.arange(1, device=device)

        torch.add(x, y, out=z)
        if device == "cuda":
            torch.cuda.synchronize()
        start = time.time()
        for i in range(repeat):
            torch.add(x, y, out=z)
        # note we deliberately do not use torch.cuda.synchronize()
        # because we want to see the overhead of the FFI call.
        end = time.time()
        print_speed(f"torch.add[{device}]", (end - start) / repeat)

    # rough take away: add on cuda roughly takes 3e-6 sec/call
    run_bench("cpu")
    run_bench("cuda")


def baseline_numpy_add(repeat):
    """Run numpy.add with one element"""
    x = np.arange(1)
    y = np.arange(1)
    z = np.arange(1)

    np.add(x, y, out=z)
    start = time.time()
    for i in range(repeat):
        np.add(x, y, out=z)
    end = time.time()
    speed = (end - start) / repeat
    print_speed("numpy.add", speed)


def baseline_cupy_add(repeat):
    """Run cupy.add with one element"""
    try:
        import cupy
    except ImportError:
        # skip if cupy is not installed
        return
    x = cupy.arange(1)
    y = cupy.arange(1)
    z = cupy.arange(1)

    cupy.add(x, y, out=z)
    start = time.time()
    for i in range(repeat):
        cupy.add(x, y, out=z)
    end = time.time()
    speed = (end - start) / repeat
    print_speed("cupy.add", speed)


def tvm_ffi_nop(repeat):
    """Overhead of tvm FFI python call via calling a NOP.

    testing.nop is defined in c++ and do nothing.
    """
    nop = tvm_ffi.get_global_func("testing.nop")
    x = tvm_ffi.from_dlpack(torch.arange(1))
    y = tvm_ffi.from_dlpack(torch.arange(1))
    z = tvm_ffi.from_dlpack(torch.arange(1))
    nop(x, y, z)
    start = time.time()
    for i in range(repeat):
        y = tvm_ffi.from_dlpack(x)
    end = time.time()
    print_speed("tvm.ffi.nop", (end - start) / repeat)


def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):
    """run dlpack conversion + tvm.ffi.nop

    Measures overhead of running dlpack for each args then invoke
    """
    nop = tvm_ffi.get_global_func("testing.nop")
    tx = tvm_ffi.from_dlpack(x)
    ty = tvm_ffi.from_dlpack(y)
    tz = tvm_ffi.from_dlpack(z)
    nop(tx, ty, tz)

    start = time.time()
    for i in range(repeat):
        tx = tvm_ffi.from_dlpack(x)
        ty = tvm_ffi.from_dlpack(y)
        tz = tvm_ffi.from_dlpack(z)
        nop(tx, ty, tz)
    end = time.time()
    print_speed(name, (end - start) / repeat)


def tvm_ffi_nop_from_torch_dlpack(repeat):
    """run dlpack conversion + tvm.ffi.nop

    Measures overhead of running dlpack for each args then invoke
    """
    x = torch.arange(1)
    y = torch.arange(1)
    z = torch.arange(1)
    bench_ffi_nop_from_dlpack("tvm.ffi.nop+from_dlpack(torch)", x, y, z, repeat)


def tvm_ffi_nop_from_numpy_dlpack(repeat):
    """run dlpack conversion + tvm.ffi.nop

    Measures overhead of running dlpack for each args then invoke
    """
    x = np.arange(1)
    y = np.arange(1)
    z = np.arange(1)
    bench_ffi_nop_from_dlpack("tvm.ffi.nop+from_dlpack(numpy)", x, y, z, repeat)


def tvm_ffi_self_dlpack_nop(repeat):
    """run dlpack conversion + tvm.ffi.nop

    Measures overhead of running dlpack for each args then invoke
    """
    x = tvm_ffi.from_dlpack(torch.arange(1))
    y = tvm_ffi.from_dlpack(torch.arange(1))
    z = tvm_ffi.from_dlpack(torch.arange(1))
    bench_ffi_nop_from_dlpack("tvm.ffi.nop+from_dlpack(tvm)", x, y, z, repeat)


def bench_ffi_nop_from_dlpack(name, x, y, z, repeat):
    """run dlpack conversion + tvm.ffi.nop

    Measures overhead of running dlpack for each args then invoke
    """
    nop = tvm_ffi.get_global_func("testing.nop")
    tx = tvm_ffi.from_dlpack(x)
    ty = tvm_ffi.from_dlpack(y)
    tz = tvm_ffi.from_dlpack(z)
    nop(tx, ty, tz)

    start = time.time()
    for i in range(repeat):
        tx = tvm_ffi.from_dlpack(x)
        ty = tvm_ffi.from_dlpack(y)
        tz = tvm_ffi.from_dlpack(z)
        nop(tx, ty, tz)
    end = time.time()
    print_speed(name, (end - start) / repeat)


def tvm_ffi_nop_from_torch_utils_to_dlpack(repeat):
    """
    Measures overhead of running dlpack for each args then invoke
    but uses the legacy torch.utils.dlpack.to_dlpack API

    This helps to measure possible implementation overhead of torch.
    """
    nop = tvm_ffi.get_global_func("testing.nop")
    x = torch.arange(1)
    y = torch.arange(1)
    z = torch.arange(1)

    tx = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x))
    ty = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(y))
    tz = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(z))
    nop(tx, ty, tz)

    start = time.time()
    for i in range(repeat):
        tx = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(x))
        ty = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(y))
        tz = tvm_ffi.from_dlpack(torch.utils.dlpack.to_dlpack(z))
        nop(tx, ty, tz)
    end = time.time()
    speed = (end - start) / repeat
    print_speed("tvm.ffi.nop+from_dlpack(torch.utils)", speed)


def bench_tvm_ffi_nop_autodlpack(name, x, y, z, repeat):
    """
    Measures overhead of running dlpack via auto convert by directly
    take torch.Tensor as inputs.
    """
    nop = tvm_ffi.get_global_func("testing.nop")
    nop(x, y, z)
    start = time.time()
    for i in range(repeat):
        nop(x, y, z)
    end = time.time()
    speed = (end - start) / repeat
    print_speed(name, speed)


def tvm_ffi_nop_autodlpack_from_torch(repeat, device="cpu"):
    """
    Measures overhead of running dlpack via auto convert by directly
    take torch.Tensor as inputs.
    """
    # use larger to ensure alignment req is met
    x = torch.arange(1, device=device)
    y = torch.arange(1, device=device)
    z = torch.arange(1, device=device)
    bench_tvm_ffi_nop_autodlpack(f"tvm.ffi.nop.autodlpack(torch[{device}])", x, y, z, repeat)


def tvm_ffi_nop_autodlpack_from_numpy(repeat):
    """
    Measures overhead of running dlpack via auto convert by directly
    take numpy.ndarray as inputs.
    """
    # use larger to ensure alignment req is met
    x = np.arange(256)
    y = np.arange(256)
    z = np.arange(256)
    bench_tvm_ffi_nop_autodlpack("tvm.ffi.nop.autodlpack(numpy)", x, y, z, repeat)


def bench_to_dlpack(x, name, repeat):
    x.__dlpack__()
    start = time.time()
    for i in range(repeat):
        x.__dlpack__()
    end = time.time()
    speed = (end - start) / repeat
    print_speed(name, speed)


def bench_to_dlpack_versioned(x, name, repeat, max_version=(1, 1)):
    """
    Measures overhead of running dlpack with latest 1.1.
    """
    try:
        x.__dlpack__(max_version=max_version)
        start = time.time()
        for i in range(repeat):
            x.__dlpack__(max_version=max_version)
        end = time.time()
        speed = (end - start) / repeat
        print_speed(name, speed)
    except Exception as e:
        print_error(name, e)


def bench_torch_utils_to_dlpack(repeat):
    """
    Measures overhead of running torch.utils.dlpack.to_dlpack
    """
    x = torch.arange(1)
    torch.utils.dlpack.to_dlpack(x)
    start = time.time()
    for i in range(repeat):
        torch.utils.dlpack.to_dlpack(x)
    end = time.time()
    speed = (end - start) / repeat
    print_speed("torch.utils.dlpack.to_dlpack", speed)


def main():
    repeat = 10000
    print("-----------------------------")
    print("Benchmark f(x, y, z) overhead")
    print("-----------------------------")
    baseline_numpy_add(repeat)
    baseline_torch_add(repeat)
    baseline_cupy_add(repeat)
    tvm_ffi_nop(repeat)
    tvm_ffi_nop_from_torch_dlpack(repeat)
    tvm_ffi_nop_from_numpy_dlpack(repeat)
    tvm_ffi_self_dlpack_nop(repeat)
    tvm_ffi_nop_from_torch_utils_to_dlpack(repeat)
    tvm_ffi_nop_autodlpack_from_torch(repeat, "cpu")
    tvm_ffi_nop_autodlpack_from_torch(repeat, "cuda")
    tvm_ffi_nop_autodlpack_from_numpy(repeat)
    print("-------------------------------")
    print("Benchmark x.__dlpack__ overhead")
    print("-------------------------------")
    bench_torch_utils_to_dlpack(repeat)
    bench_to_dlpack(torch.arange(1), "torch.__dlpack__", repeat)
    bench_to_dlpack(np.arange(1), "numpy.__dlpack__", repeat)
    bench_to_dlpack(tvm_ffi.from_dlpack(torch.arange(1)), "tvm.__dlpack__", repeat)
    print("---------------------------------------------------")
    print("Benchmark x.__dlpack__(max_version=(1,1)) overhead")
    print("---------------------------------------------------")
    bench_to_dlpack_versioned(torch.arange(1), "torch.__dlpack__(max_version=(1,1))", repeat)
    bench_to_dlpack_versioned(np.arange(1), "numpy.__dlpack__(max_version=(1,1))", repeat)
    bench_to_dlpack_versioned(
        tvm_ffi.from_dlpack(torch.arange(1)), "tvm.__dlpack__(max_version=(1,1))", repeat
    )


if __name__ == "__main__":
    main()