# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import time
import sys
import torch

def measure_cpu(a, b, steps):

    global c
    start = time.perf_counter()
    for i in range(steps):
        c = torch.mm(a, b)
    end = time.perf_counter()
    c.to('cpu')
    return end - start


def measure_gpu(a, b, steps):

    global c
    torch.cuda.synchronize()
    start = time.perf_counter()
    for i in range(steps):
        c = torch.mm(a, b)
    torch.cuda.synchronize()
    end = time.perf_counter()
    c.to('cpu')
    return end - start


def measure_xla(a, b, steps):

    import torch_xla

    def sync(tensor, dev):
        torch_xla._XLAC._xla_sync_multi([tensor], devices=[str(dev)], wait=True, sync_xla_data=True)

    c = torch.mm(a, b)

    start = time.perf_counter()
    for _ in range(steps):
        # Add data dependency to prevent loop elimination
        # The PyTorch/XLA lazy evaluation will eliminate the loop
        # Simplier data dependency will not work
        b[0] = torch.min(c[0], b[0])
        c = torch.min(c, torch.mm(a, b))

    sync(c, c.device)
    end = time.perf_counter()
    # c.to('cpu')
    return end - start

def run_single(args, m, n, k):

    dtype = args.dtype
    device = args.device
    warmups = args.warmups
    steps = args.steps

    dt = torch.float32
    if (dtype == "float16" or dtype == "half"):
        dt = torch.float16
    elif (dtype == "bfloat16"):
        dt = torch.bfloat16

    torch.manual_seed(0)

    elap = 0.0

    a = torch.randn(m, k).to(dt)
    b = torch.randn(k, n).to(dt)
    c = torch.zeros(m, n).to(dt)

    if device == 'cpu':

        measure_cpu(a, b, warmups)
        elap = measure_cpu(a, b, steps)

    elif device == 'gpu':

        if torch.cuda.is_available():
            # ncuda = torch.cuda.device_count()
            # print("There are {} cuda devices".format(ncuda))
            # print("The first cuda device name is {} ".format(torch.cuda.get_device_name()))
            cuda0 = torch.device('cuda:0')
            with torch.cuda.device(cuda0):
                acuda = a.to(cuda0)
                bcuda = b.to(cuda0)
                measure_gpu(acuda, bcuda, warmups)
                elap = measure_gpu(acuda, bcuda, steps)
        else:
            print("CUDA is not available")
            sys.exit(1)

    else:
        # import torch_xla
        import torch_xla.core.xla_model as xm

        # alldev = xm.get_xla_supported_devices()
        # allrealdev = xm.xla_real_devices(alldev)
        # print("Found {0} XLA devices: {1}".format(len(allrealdev), allrealdev))

        dev = xm.xla_device()
        a = a.to(dev)
        b = b.to(dev)
        c = c.to(dev)
        measure_xla(a, b, warmups)
        xm.mark_step()
        elap = measure_xla(a, b, steps)
        xm.mark_step()

    return elap

def run(args, dataset):

    print("----------------------------------------------------------------")
    print("         M         N          K          Time(s)      Rate(TF/s)")
    print("----------------------------------------------------------------")
    for i in range(len(dataset)):
        m, n, k = dataset[i]
        elap = run_single(args, m, n, k)
        elap /= args.steps
        print("{0:10}, {1:10}, {2:10},     {3:10.6f}     {4:.3f} ".format(m, n, k, elap,
            m * n * k * 2 * 1.0 / elap / 1.0e12))

if __name__ == "__main__":

    import argparse

    parser = argparse.ArgumentParser(
        description="Measure the performance of GEMM using mm, or matmul")
    # model related parameters
    parser.add_argument("-m", "--msize", type=int, default=1024)
    parser.add_argument("-n", "--nsize", type=int, default=1024)
    parser.add_argument("-k", "--ksize", type=int, default=1024)
    parser.add_argument("-t", "--dtype", type=str, default="float32")
    parser.add_argument("-d", "--device", choices=['cpu', 'gpu', 'tpu'], type=str, default='cpu')
    parser.add_argument("--steps", type=int, default=100)
    parser.add_argument("--warmups", type=int, default=10)
    args = parser.parse_args()

    d = [(args.msize, args.nsize, args.ksize)]
    run(args, d)

