build/torch27-cxx11-cu126-x86_64-linux/megablocks/benchmark_util.py

# Copyright 2024 Databricks
# SPDX-License-Identifier: Apache-2.0

import numpy as np
import torch


def log_benchmark(name, arguments, time, std):
    print('=' * 60)
    print(f'{name} Benchmark')
    print('Benchmark Parameters:')
    for (key, value) in arguments.items():
        print(f'{key} = {value}')
    print('Results:')
    print('mean time = {:.3f}ms, std time = {:.3f}ms'.format(time, std))
    print('=' * 60)


def benchmark_function(fn, iterations=100, warmup=10):
    # Warmup iterations.
    for _ in range(warmup):
        fn()

    times = []
    for i in range(iterations):
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)

        start.record()
        fn()
        end.record()

        torch.cuda.synchronize()
        times.append(start.elapsed_time(end))
    return np.mean(times), np.std(times)