cl math test

clmath_test.py

import numpy as np

import pyopencl as cl
from pyopencl import tools as cl_tools
from pyopencl import array as cl_array
from pyopencl import clmath

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
mf = cl.mem_flags
mem_pool = cl_tools.MemoryPool(
    cl_tools.ImmediateAllocator(queue, mf.READ_WRITE))

a = np.complex64(np.random.randn(1000, 64) + 1j*np.random.randn(1000, 64))

cl_a_real = cl_array.Array(queue, a.shape, np.float32, allocator=mem_pool)
cl_a_imag = cl_array.Array(queue, a.shape, np.float32, allocator=mem_pool)

cl_a_real.set(a.real.copy())
cl_a_imag.set(a.imag.copy())

def cl_test():
    for n in range(10000):
        foo = clmath.atan2(cl_a_imag, cl_a_real, queue=queue)

    cl.enqueue_barrier(queue).wait()

cl_test_out = cl_array.Array(queue, a.shape, np.float32, allocator=mem_pool)

def cl_test_preallocated():
    for n in range(10000):
        clmath._atan2(cl_test_out, cl_a_imag, cl_a_real, queue=queue)
    
    cl.enqueue_barrier(queue).wait()

import cProfile
cProfile.run('cl_test()', 'cl_test_pypi')
cProfile.run('cl_test_preallocated()', 'cl_test_preallocated_pypi')