import time
import torch

B, C, H, W = 8, 3, 1080, 1920
device = "cuda"

print("torch version:", torch.__version__)
print("device:", device)

imgs = torch.randint(
    0, 256, (B, C, H, W),
    dtype=torch.uint8,
    device="cpu",
)
imgs_pinned = imgs.pin_memory()
cpu_pinned_float = torch.empty((B, C, H, W), dtype=torch.float32, device="cpu", pin_memory=True)
gpu_float = torch.empty((B, C, H, W), dtype=torch.float32, device=device)


def benchmark(name, fn, iters=50):
    fn()
    torch.cuda.synchronize()

    t0 = time.perf_counter()
    for _ in range(iters):
        fn()
        torch.cuda.synchronize()
    t1 = time.perf_counter()

    print(f"{name:35s}: {(t1 - t0) * 1000 / iters:.3f} ms/iter")


def cpu_uint8_to_float32():
    return imgs.to(dtype=torch.float32)


def cpu_uint8_to_cuda_float32_combined():
    return imgs.to(device=device, dtype=torch.float32).div_(255.0)


def cpu_uint8_to_cuda_then_dtype():
    return imgs.to(device=device).to(dtype=torch.float32).div_(255.0)


def cpu_uint8_to_float32_then_cuda():
    x = imgs.to(dtype=torch.float32)
    return x.to(device=device).div_(255.0)


def cpu_uint8_pinned_to_cuda_float32_combined_nonblocking():
    return imgs_pinned.to(device=device, dtype=torch.float32, non_blocking=True).div_(255.0)


def cpu_uint8_pinned_to_cuda_then_dtype_nonblocking():
    return imgs_pinned.to(device=device, non_blocking=True).to(dtype=torch.float32).div_(255.0)


def cpu_uint8_copy_to_pinned_float_then_gpu_nonblocking():
    cpu_pinned_float.copy_(imgs)
    gpu_float.copy_(cpu_pinned_float, non_blocking=True)
    return gpu_float


if __name__ == "__main__":
    benchmark("cpu uint8 -> float32", cpu_uint8_to_float32)
    benchmark("cpu uint8 -> cuda float32 (combined)", cpu_uint8_to_cuda_float32_combined)
    benchmark("cpu uint8 -> cuda -> float32", cpu_uint8_to_cuda_then_dtype)
    benchmark("cpu uint8 -> float32 -> cuda", cpu_uint8_to_float32_then_cuda)
    benchmark("cpu uint8 pinned -> cuda float32 (combined, nb)", cpu_uint8_pinned_to_cuda_float32_combined_nonblocking)
    benchmark("cpu uint8 pinned -> cuda -> float32 (nb)", cpu_uint8_pinned_to_cuda_then_dtype_nonblocking)
    benchmark("cpu uint8 -> pinned float32 -> cuda (nb, copy_)", cpu_uint8_copy_to_pinned_float_then_gpu_nonblocking)