I have been trying to execute a split-step method to numerically integrate the Gross-Pitaevskii equation. My code was performing as expected using python, but in order to get a performance increase I have been adapting it using PyOpenCL, so that it can run on a GPU. I seem to have gotten it working, as it agrees with the results my python code (running on the CPU) is giving, but it takes much more time than I expected. A working example can be found below:
###########################################################################
# IMPORTS NECESSARY TO RUN
from __future__ import absolute_import
from __future__ import print_function
import numpy as np
import scipy.fftpack as sp
import time as time
import pyopencl as cl
import pyopencl.array as cl_array
from reikna.cluda import ocl_api
from reikna.core import Computation, Parameter, Annotation
from reikna.fft import FFT, FFTShift
################################################################################
# DEFINE SYSTEM PARAMETERS
Lx = 1000 # length of system in x-direction
Ly = 500 # length of system in y-direction
dx = 0.4 # space step
dt = 0.1*(dx**2) # calculated timestep
Nx = int(Lx/dx) # number of points in x-direction
Ny = int(Ly/dx) # number of points in y-direction
# create frequency space coordinate grid
kx = np.array([-Nx*np.pi/Lx+2.0*np.pi*i/Lx for i in range(Nx)]) #x-wavenumbers
# ^ used when we Fourier transform
ky = np.array([-Ny*np.pi/Ly+2.0*np.pi*i/Ly for i in range(Ny)]) #x-wavenumbers
kxx, kyy = np.meshgrid(kx, ky, sparse=True) #wavenumbergrid
#################################################################################
# GENERATE TRAP POTENTIAL AND INITIAL VALUES
# define the trap potential matrix (constant)
# it has a value of 100 at the edge, and zero in the bulk
Vmat = 100.0*np.ones((Ny, Nx))
Vmat[4:-4,4:-4] = np.zeros((Ny - 8, Nx - 8))
# the initial wavefunction is zero at the edge (where the potential is nonzero)
# and 1 in the bulk (where the potential is zero).
U0 = np.zeros((Ny, Nx))
U0[4:-4,4:-4] = np.ones((Ny - 8, Nx - 8))
U = U0
###################################################################################
# PASS ARRAYS TO DEVICE
# define the PyOpenCL queue
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)
# define the Reikna thread
api = ocl_api()
thr = api.Thread(queue)
# make sure all of the data types are correct
U = U.astype(np.complex128)
Vmat = Vmat.astype(np.float64)
L_op = np.exp(-1j*dt*0.5*( kxx**2+kyy**2 )).astype(np.complex128)
idtmat = np.array(1j*dt).astype(np.complex128) # we will use the 1j below
# pass the arrays to the device, all using pyopencl, can later use w/ reikna
op_dev = cl_array.to_device(queue, U)
Vmat_dev = cl_array.to_device(queue, Vmat)
L_op_dev = cl_array.to_device(queue, L_op)
idt_dev = cl_array.to_device(queue, idtmat)
###################################################################################
# PYOPENCL KERNEL DEFINITIONS
gpe = cl.Program(ctx, """
#pragma OPENCL EXTENSION cl_khr_fp64: enable
#define PYOPENCL_DEFINE_CDOUBLE
#include <pyopencl-complex.h>
__kernel void nonlinear_operator(__global const cdouble_t *U, __global const double *Vmat, __global const cdouble_t *idt, __global cdouble_t *U_aft)
{
// get thread id numbers
int gid0 = get_global_id(0);
int gid1 = get_global_id(1);
// get the global size of the blocks
int num0 = get_global_size(0);
// the real value that gets exponentiated
__local double mag;
mag = cdouble_abs(U[gid0 + num0 * gid1]) * cdouble_abs(U[gid0 + num0 * gid1]) + Vmat[gid0 + num0 * gid1];
// exponentiate and multiply real exponent by complex wavefct
U_aft[gid0 + num0 * gid1] = cdouble_mul( cdouble_exp( cdouble_mulr(idt[0], -mag) ), U[gid0 + num0 * gid1]);
}
__kernel void laplacian_operator(__global const cdouble_t *C, __global const cdouble_t *L_op, __global cdouble_t *C_aft)
{
// get thread id numbers
int gid0 = get_global_id(0);
int gid1 = get_global_id(1);
// get the global sizes of the various blocks
int num0 = get_global_size(0);
// exponentiate and multiply real exponent by complex wavefct
C_aft[gid0 + num0 * gid1] = cdouble_mul( L_op[gid0 + num0 * gid1], C[gid0 + num0 * gid1]);
}
""").build()
###################################################################################
# REIKNA KERNEL DEFINITIONS
fft = FFT(U)
fftc = fft.compile(thr)
shift = FFTShift(U)
shiftc = shift.compile(thr)
##############################################################################
# RUNNING THE KERNELS, TIMING INCLUDED
t0 = time.time()
t_loop = []
for i in range(100):
t0_loop = time.time()
# apply the nonlinear operator
gpe.nonlinear_operator(queue, op_dev.shape, None, op_dev.data, Vmat_dev.data, idt_dev.data, op_dev.data)
# transform to frequency space
fftc(op_dev, op_dev)
# apply the shift operator to get zero frequency components to center
shiftc(op_dev, op_dev)
# apply the Laplacian operator in frequency space
gpe.laplacian_operator(queue, op_dev.shape, None, op_dev.data, L_op_dev.data, op_dev.data)
# shift back
shiftc(op_dev, op_dev)
# transform back to position space
fftc(op_dev, op_dev, inverse=True)
t_loop.append(time.time() - t0_loop)
Delta_t = time.time()-t0
# Copy the array back from the device
t_copy = time.time()
final_array = op_dev.get()
Delta_tcopy = time.time()-t_copy
##################################################################################
# COMPARE TO CALCULATION DONE ON CPU
t1 = time.time()
cpu_U = U
for i in range(100):
cpu_U=np.exp( -1j*dt*( Vmat + cpu_U*np.conjugate(cpu_U) ))*cpu_U #advance w/ N op
cpu_C=sp.fftshift(sp.fft2(cpu_U)) # transform to fourier space
cpu_C=np.exp( -1j*dt*0.5*( kxx**2+kyy**2 ) )*cpu_C # advance w/ L op
cpu_U=sp.ifft2(sp.fftshift(cpu_C)) # transform back
Delta_t1 = time.time() - t1
test = np.amax(abs(final_array-cpu_U))
if test <= 10**(-6.0):
print('Correct!')
print('GPU takes '+str(Delta_t)+' sec.')
print('Copying takes '+str(Delta_tcopy)+' sec.')
print('CPU Python takes '+str(Delta_t1)+' sec.')
else:
print('Not correct!')
print(test)
################################################################################
# WRITE OUT THE INDIVIDUAL LOOP TIMES INTO A FILE
target = open('loop_times.txt','w')
for i in range(len(t_loop)):
target.write('Loop number '+str(i)+' takes '+str(t_loop[i])+' seconds to complete.')
target.write('\n')
target.close()
If this code is run it shows that the CPU and GPU results agree. However, running it on a Tesla K40c reveals that the GPU only runs about 10x faster than the CPU. Inspecting the loop_times.txt file, which contains the timing information for each loop on the GPU, shows that the loops are initially very fast. Despite this initial speed, after about 20 loops they become 200x slower than they were before, and stay at that speed for the rest of the computation. Does anyone have any ideas why this could be? My best guess is that it's a problem with the memory. The PyOpenCL documentation here states that "... code based on pyopencl.array.Array can easily run into issue[s] because a fresh memory area is allocated for each intermediate result". I'm not positive this is the issue, though, since I don't declare intermediate arrays.
I have included the loop_times.txt file from a run on a Tesla K40c below, in case this is helpful for diagnosing the problem. Thanks in advance!
Loop number 0 takes 0.00145196914673 seconds to complete.
Loop number 1 takes 0.000530004501343 seconds to complete.
Loop number 2 takes 0.000539064407349 seconds to complete.
Loop number 3 takes 0.000540018081665 seconds to complete.
Loop number 4 takes 0.000539064407349 seconds to complete.
Loop number 5 takes 0.00052809715271 seconds to complete.
Loop number 6 takes 0.000566959381104 seconds to complete.
Loop number 7 takes 0.000523090362549 seconds to complete.
Loop number 8 takes 0.000649929046631 seconds to complete.
Loop number 9 takes 0.000531196594238 seconds to complete.
Loop number 10 takes 0.000524997711182 seconds to complete.
Loop number 11 takes 0.000524997711182 seconds to complete.
Loop number 12 takes 0.000520944595337 seconds to complete.
Loop number 13 takes 0.000530004501343 seconds to complete.
Loop number 14 takes 0.000522136688232 seconds to complete.
Loop number 15 takes 0.000525951385498 seconds to complete.
Loop number 16 takes 0.000523090362549 seconds to complete.
Loop number 17 takes 0.0888669490814 seconds to complete.
Loop number 18 takes 0.102005958557 seconds to complete.
Loop number 19 takes 0.102015972137 seconds to complete.
Loop number 20 takes 0.102032899857 seconds to complete.
Loop number 21 takes 0.101969957352 seconds to complete.
Loop number 22 takes 0.102008104324 seconds to complete.
Loop number 23 takes 0.102007865906 seconds to complete.
Loop number 24 takes 0.10200715065 seconds to complete.
Loop number 25 takes 0.102005004883 seconds to complete.
Loop number 26 takes 0.102000951767 seconds to complete.
Loop number 27 takes 0.102005004883 seconds to complete.
Loop number 28 takes 0.102003097534 seconds to complete.
Loop number 29 takes 0.101999998093 seconds to complete.
Loop number 30 takes 0.101995944977 seconds to complete.
Loop number 31 takes 0.101994037628 seconds to complete.
Loop number 32 takes 0.10199713707 seconds to complete.
Loop number 33 takes 0.101987838745 seconds to complete.
Loop number 34 takes 0.102010011673 seconds to complete.
Loop number 35 takes 0.102000951767 seconds to complete.
Loop number 36 takes 0.102009057999 seconds to complete.
Loop number 37 takes 0.102005004883 seconds to complete.
Loop number 38 takes 0.102013111115 seconds to complete.
Loop number 39 takes 0.102020025253 seconds to complete.
Loop number 40 takes 0.102008104324 seconds to complete.
Loop number 41 takes 0.102012872696 seconds to complete.
Loop number 42 takes 0.102003097534 seconds to complete.
Loop number 43 takes 0.102008104324 seconds to complete.
Loop number 44 takes 0.101997852325 seconds to complete.
Loop number 45 takes 0.102009773254 seconds to complete.
Loop number 46 takes 0.102011919022 seconds to complete.
Loop number 47 takes 0.101995944977 seconds to complete.
Loop number 48 takes 0.102001905441 seconds to complete.
Loop number 49 takes 0.102009057999 seconds to complete.
Loop number 50 takes 0.101994037628 seconds to complete.
Loop number 51 takes 0.102015018463 seconds to complete.
Loop number 52 takes 0.10200715065 seconds to complete.
Loop number 53 takes 0.102021932602 seconds to complete.
Loop number 54 takes 0.102017879486 seconds to complete.
Loop number 55 takes 0.102023124695 seconds to complete.
Loop number 56 takes 0.102003097534 seconds to complete.
Loop number 57 takes 0.102006912231 seconds to complete.
Loop number 58 takes 0.10199713707 seconds to complete.
Loop number 59 takes 0.102031946182 seconds to complete.
Loop number 60 takes 0.102022171021 seconds to complete.
Loop number 61 takes 0.102020025253 seconds to complete.
Loop number 62 takes 0.102014064789 seconds to complete.
Loop number 63 takes 0.102007865906 seconds to complete.
Loop number 64 takes 0.101998090744 seconds to complete.
Loop number 65 takes 0.102015018463 seconds to complete.
Loop number 66 takes 0.102014064789 seconds to complete.
Loop number 67 takes 0.102025032043 seconds to complete.
Loop number 68 takes 0.102019071579 seconds to complete.
Loop number 69 takes 0.102022886276 seconds to complete.
Loop number 70 takes 0.102005958557 seconds to complete.
Loop number 71 takes 0.102015972137 seconds to complete.
Loop number 72 takes 0.102024078369 seconds to complete.
Loop number 73 takes 0.101996898651 seconds to complete.
Loop number 74 takes 0.102014064789 seconds to complete.
Loop number 75 takes 0.10201215744 seconds to complete.
Loop number 76 takes 0.102012872696 seconds to complete.
Loop number 77 takes 0.101979017258 seconds to complete.
Loop number 78 takes 0.101991176605 seconds to complete.
Loop number 79 takes 0.102010011673 seconds to complete.
Loop number 80 takes 0.102005958557 seconds to complete.
Loop number 81 takes 0.102019071579 seconds to complete.
Loop number 82 takes 0.102010965347 seconds to complete.
Loop number 83 takes 0.102006912231 seconds to complete.
Loop number 84 takes 0.101999044418 seconds to complete.
Loop number 85 takes 0.102009057999 seconds to complete.
Loop number 86 takes 0.102022886276 seconds to complete.
Loop number 87 takes 0.10201382637 seconds to complete.
Loop number 88 takes 0.101995944977 seconds to complete.
Loop number 89 takes 0.102017879486 seconds to complete.
Loop number 90 takes 0.102014064789 seconds to complete.
Loop number 91 takes 0.10200214386 seconds to complete.
Loop number 92 takes 0.101999998093 seconds to complete.
Loop number 93 takes 0.102025032043 seconds to complete.
Loop number 94 takes 0.102019071579 seconds to complete.
Loop number 95 takes 0.101996898651 seconds to complete.
Loop number 96 takes 0.102020025253 seconds to complete.
Loop number 97 takes 0.101989984512 seconds to complete.
Loop number 98 takes 0.102004051208 seconds to complete.
Loop number 99 takes 0.102003097534 seconds to complete.
Since you do not synchronize the queue at the end of each iteration, what you are measuring is essentially the enqueuing time. It seems that by 17th iteration the queuing buffer is full, and each new kernel call has to wait until another kernel finishes and gets removed from the queue. This gives you roughly correct times starting from that moment. When I run it on my machine (iMac with gf755m), the queue actually accepts all 100 iterations of kernels, and then I have to wait for a minute for it to churn through all of them, giving me a result like
GPU takes 0.055264949798583984 sec.
Copying takes 52.194445848464966 sec.
If you want to measure the actual iteration time, you need to synchronize the queue at the end of each iteration:
queue.finish()
or, equivalently, through the Reikna API
thr.synchronize()
There are several things worth noting.
I got some strange address boundary errors when I ran your code. This seems to be caused by the fact that PyOpenCL's complex value handling functions like cdouble_mul are actually macros, and the error appears when you give them AST bits that access the memory, for instance
cdouble_mul(L_op[gid0 + num0 * gid1], C[gid0 + num0 * gid1])
The error disappears if you pre-load the values before passing them to the macro:
cdouble_t L_op_val = L_op[gid0 + num0 * gid1];
cdouble_t C_val = C[gid0 + num0 * gid1];
cdouble_mul(L_op_val, C_val);
You don't need to create a 1-element array out of 1j*dt. You can just pass it to the kernel as a scalar.
Perhaps you are aware of that, but just in case: you pass the array shape as global size to your kernels, so inside them you are essentially addressing an NxM array as an MxN array. As long as the array is contiguous in memory and all your operations are elementwise, it doesn't matter, but make sure you keep that in mind.
If I write the nonlinear_operator kernel using Reikna as
gpe2 = reikna.algorithms.PureParallel([
Parameter('U', Annotation(U, 'i')),
Parameter('Vmat', Annotation(Vmat, 'i')),
Parameter('idt', Annotation(idtmat.reshape(1)[0], 's')),
Parameter('U_aft', Annotation(U, 'o'))],
"""
${U.ctype} U = ${U.load_same};
${Vmat.ctype} U_squared = ${norm}(U);
${Vmat.ctype} mag = U_squared + ${Vmat.load_same};
${U_aft.store_same}(
${mul_cc}(
${exp}(${mul_cr}(${idt}, -mag)),
U));
""",
render_kwds=dict(
norm=reikna.cluda.functions.norm(U.dtype),
mul_cr=reikna.cluda.functions.mul(U.dtype, Vmat.dtype),
mul_cc=reikna.cluda.functions.mul(U.dtype, U.dtype),
exp=reikna.cluda.functions.exp(U.dtype)))
gpe2c = gpe2.compile(thr)
# in the loop
gpe2c(op_dev, Vmat_dev, idtmat[0], op_dev)
it works twice as fast as your original kernel. It doesn't matter that much, of course, because the majority of time is spent in the FFT.
Again, you may be aware of that, but Reikna's FFT for non-power-of-2 problem sizes uses the Bluestein's algorithm which essentially doubles the work. Also, instead of doing two fftshifts on a GPU in each iteration, it will be faster to shift L_op once before the loop (you can also generate it using numpy.fft.fftfreq which gives you the correct order of frequencies right away).