Function timed

tokasaurus/utils.py:200–251

(
    fn,
    num_iters=50,
    num_warmup=10,
    unit: str = "ms",
    between_fn=None,
    prog=True,
    barrier=False,
)

@torch.no_grad()
def timed(
    fn,
    num_iters=50,
    num_warmup=10,
    unit: str = "ms",
    between_fn=None,
    prog=True,
    barrier=False,
):
    warmup_times = []
    times = []
    cpu_warmup_times = []
    cpu_times = []
    for itr in tqdm(range(num_iters + num_warmup), desc="Timing", disable=not prog):
        if between_fn is not None:
            between_fn()
            torch.cuda.synchronize()

        if barrier:
            torch.distributed.barrier()

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()  # type: ignore
        cpu_start = time.perf_counter()
        _ = fn()
        cpu_end = time.perf_counter()
        end.record()  # type: ignore
        torch.cuda.synchronize()

        if barrier:
            torch.distributed.barrier()

        gpu_milis = start.elapsed_time(end)
        cpu_milis = (cpu_end - cpu_start) * 1000

        gpu_time = convert_unit(gpu_milis, unit)
        cpu_time = convert_unit(cpu_milis, unit)

        if itr >= num_warmup:
            times.append(gpu_time)
            cpu_times.append(cpu_time)
        else:
            warmup_times.append(gpu_time)
            cpu_warmup_times.append(cpu_time)

    return TimeResult(
        times=times,
        warmup_times=warmup_times,
        cpu_times=cpu_times,
        cpu_warmup_times=cpu_warmup_times,
    )


@torch.no_grad()