from __future__ import absolute_import, division, print_function, unicode_literals
import time
from collections import namedtuple
from torch.utils import ThroughputBenchmark

NUM_LOOP_ITERS = 1000
BenchmarkConfig = namedtuple('BenchmarkConfig', 'num_warmup_iters num_iters')
ModuleConfig = namedtuple('ModuleConfig', 'pt_fn c2_op num_params graph_mode')

def ms_to_us(time_ms):
    return (time_ms * 1e3)

def secs_to_us(time_s):
    return (time_s * 1e6)

def secs_to_ms(time_s):
    return (time_s * 1e3)

def benchmark_using_throughput_benchmark(config, module):
    print("Benchmarking via ThroughputBenchmark")
    bench = ThroughputBenchmark(module.module)
    bench.add_input(*module.tensor_inputs)
    stats = bench.benchmark(1, config.num_warmup_iters, config.num_iters)
    return stats.latency_avg_ms / NUM_LOOP_ITERS

def benchmark_module(config, module, use_throughput_benchmark=False):
    if use_throughput_benchmark:
        return benchmark_using_throughput_benchmark(config, module)
    module.forward(config.num_warmup_iters)
    print("Running module for {} iterations".format(config.num_iters))
    start = time.time()
    module.forward(config.num_iters)
    end = time.time()
    time_elapsed_s = (end - start)
    return (secs_to_ms(time_elapsed_s) / config.num_iters / NUM_LOOP_ITERS)