import time
import argparse
import torch

parser = argparse.ArgumentParser()
parser.add_argument("--events", action="store_true", help="Use CUDA events")
parser.add_argument("--check", action="store_true", help="Enable correctness check")
args = parser.parse_args()

shapes = [[(34816, 1), (34816, 3840)]]


for i, shape in enumerate(shapes):
    a = torch.randn(shape[0], device='cuda', dtype=torch.float)
    b = torch.randn(shape[1], device='cuda', dtype=torch.bfloat16)
    for i in range(20):
        if args.check and i == 5:
            a_cpu = a.cpu()
            b_cpu = b.cpu()
            c_cpu = torch.mul(a_cpu, b_cpu)
            c = torch.mul(a, b)
            torch.cuda.synchronize()
            assert torch.equal(c.cpu(), c_cpu)
        _ = torch.mul(a, b)
    torch.cuda.synchronize()

    if args.events:
        start_evt = torch.cuda.Event(enable_timing=True)
        end_evt = torch.cuda.Event(enable_timing=True)
        start_evt.record()
    else:
        start_time = time.perf_counter_ns()

    for _ in range(100):
        c = torch.mul(a, b)

    if args.events:
        end_evt.record()
    else:
         torch.cuda.synchronize()
         end_time = time.perf_counter_ns()

    if args.events:
        torch.cuda.synchronize()

    if args.events:
        print(f"Avg time for shape {shape}: {start_evt.elapsed_time(end_evt) / 100 * 1e3:.2f} us")
    else:
        print(f"Avg time for shape {shape}: {(end_time - start_time) / (100 * 1e3):.2f} us")

Before:
Avg time for shape [(34816, 1), (34816, 3840)]: 406.78 us

After:
Avg time for shape [(34816, 1), (34816, 3840)]: 231.73 us