PyOpenCL bandwidth measurement

bandwidth.py

import numpy as np
import pyopencl as cl
from time import time


def bandwidth_calculator(n_numbers):
    a = np.random.rand(n_numbers).astype(np.float32)

    mf = cl.mem_flags
    a_dev = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=a)
    a_twice_dev = cl.Buffer(ctx, mf.WRITE_ONLY, a.nbytes)

    prg = cl.Program(ctx, """
        __kernel void twice(
                __global const float *a_dev, __global float *a_twice_dev)
        {
            a_twice_dev[get_global_id(0)] = a_dev[get_global_id(0)] * 2;
        }
        """).build()

    knl = prg.twice
    a_twice = np.empty_like(a)

    for i in range(3):
        evt = knl(queue, a.shape, None, a_dev, a_twice_dev)

    nruns = 5
    queue.finish()
    gpu_start_time = time()

    for i in range(nruns):
        evt = knl(queue, a.shape, None, a_dev, a_twice_dev)

    evt.wait()

    gpu_end_time = time()

    cl.enqueue_copy(queue,  a_twice, a_twice_dev)
    return (nruns*2*a.nbytes)/((gpu_end_time-gpu_start_time)*1e9)


if __name__ == "__main__":
    ctx = cl.create_some_context()
    queue = cl.CommandQueue(ctx)

    number_list = []
    bandwidth_list = []
    for k in range(20, 29):
        n = 2**k
        number_list.append(n)
        bandwidth_list.append(bandwidth_calculator(n))
        print(
                "Number of floats=%s\tBandwidth obtained = %.1f GB/s"
                % (n, bandwidth_list[-1]))