synctest.py

# Example output
"""
f            : 0.000139 [s]
sync after f : 0.870124 [s]

f2            : 0.820720 [s]
sync after f2 : 0.000009 [s]

g            : 0.821427 [s]
sync after g : 0.000004 [s]
"""

import time
import chainer
import cupy

N = 2048

dev = cupy.cuda.Device()
pool = cupy.cuda.MemoryPool()
handle = cupy.cuda.cublas.create()

A = pool.malloc(8 * N * N)
B = pool.malloc(8 * N * N)
C = pool.malloc(8 * N * N)

def f(rep=10):
    for i in range(rep):
        cupy.cuda.cublas.dgemm(handle, 0, 0, N, N, N, 1, A, N, B, N, 1, C, N)

def f2(rep=10):
    f(rep)
    dev.synchronize()

def g(rep=10):
    for i in range(rep):
        cupy.cuda.cublas.dgemm(handle, 0, 0, N, N, N, 1, A, N, B, N, 1, C, N)
        dev.synchronize()


def measure(func):
    beg1 = time.time()
    func()
    end1 = time.time()
        
    beg2 = time.time()
    dev.synchronize()  # sync anyways
    end2 = time.time()

    t1 = (end1 - beg1)
    t2 = (end2 - beg2)
    print("{}            : {:.6f} [s]".format(func.__name__, t1))
    print("sync after {} : {:.6f} [s]".format(func.__name__, t2))
    print()
    
def main():
    measure(f)
    measure(f2)
    measure(g)


if __name__ == '__main__':
    main()