import torch
from time import time
import os

N = 128
T = 30
H = 500
count = 1000

def bench_bmm():
    mat1 = torch.randn(N, 1, T)
    mat2 = torch.randn(N, T, H)

    tstart = time()
    for i in range(count):
        res = torch.bmm(mat1, mat2)
    tend = time()
    t = (tend-tstart)/count*1000
    flops = N*1*T*H*2 / t / 1000000
    print("run torch.bmm:")
    print("total time     : %.2f s" % (tend-tstart))
    print("each iteration : %.2f ms %.2f GFlops" % (t, flops))

def test_bmm(trans_A=False, trans_B=False):
    I = 10
    print("testing torch.bmm mat1: %s, mat2 %s" % ("T" if trans_A else "N", "T" if trans_B else "N"))
    mat1 = torch.randn(N, T, I).transpose(1, 2) if trans_A else torch.randn(N, I, T)
    mat2 = torch.randn(N, H, T).transpose(1, 2) if trans_B else torch.randn(N, T, H)
    mat1_ = mat1.clone()
    mat2_ = mat2.clone()

    res = torch.bmm(mat1, mat2)
    res_ = torch.Tensor(N, I, H)
    for i in range(N):
        res_[i] = mat1_[i].matmul(mat2_[i])

    for ii in range(N):
        for jj in range(I):
            for kk in range(H):
                val1 = res[ii][jj][kk]
                val2 = res_[ii][jj][kk]
                if abs(val1-val2) < 1e-5:
                    continue
                else:
                    print(res[ii][jj][kk], res_[ii][jj][kk], "not equal, case FAIL")
                    return

    print("PASS")

bench_bmm()
test_bmm(trans_A=False, trans_B=False)
test_bmm(trans_A=True, trans_B=False)
test_bmm(trans_A=False, trans_B=True)
test_bmm(trans_A=True, trans_B=True)