Measure GEMM speed via numpy.dot

dot_flops.py

#!/bin/python
import sys
import time
import numpy as np

N = int(sys.argv[1])
dtype = sys.argv[2]

A = np.random.rand(N, N).astype(dtype)
B = np.random.rand(N, N).astype(dtype)

t0 = time.time()
reps = 5
for ii in range(reps):
    C = np.dot(A, B)
t1 = time.time()

FLOPS = reps * 2 * N ** 3
print 'GFLOP/s', FLOPS / (1000 ** 3) / (t1 - t0)