bench-dispatch.py

import torch
import sys
import timeit

def add_repeat(x, y, repeat=10):
    rc = x.clone()
    for i in range(repeat):
       rc += y
    return rc

def main():
    from torch.utils.benchmark import Timer, Compare
    measurements = []
    for repeat in [10, 100, 1000]:
        x, y = torch.rand(2, 4096, device='cuda').unbind()
        kwargs = {"x": x, "y": y, "add_repeat": add_repeat}
        stmt=f"add_repeat(x, y, repeat={repeat})"
        timer = Timer(stmt=stmt, globals=kwargs, timer=timeit.default_timer, description="time")
        measurements.append(timer.blocked_autorange())
    Compare(measurements).print()

if __name__ == "__main__":
    print(f"Python version {sys.version} Torch version {torch.__version__} GPU is {torch.cuda.get_device_name(0)}\n")
    main()