Source

yt.opencl / simple_grid_sum.py

import matplotlib;matplotlib.use("Agg");import pylab
from yt.config import ytcfg
from yt.mods import *
import pyopencl as cl
import pyopencl.array as cl_array
from contextlib import contextmanager

times = {}

@contextmanager
def time_func(name):
    t1 = time.time()
    yield
    t2 = time.time()
    print "%s took %0.3e seconds" % (name, t2-t1)
    times[name] = t2-t1

pf = load("JHK-DD0030/galaxy0030")
#pf = load("DD0087/DD0087")

# This part is specific to Hex
platforms = cl.get_platforms()
GPU = platforms[0]
CPU = platforms[1]
devices = GPU.get_devices()

gformat = cl.ImageFormat(cl.channel_order.R,
                         cl.channel_type.FLOAT)

context = cl.Context(devices)
command_queue = cl.CommandQueue(context)

with time_func("Load from disk"):
    for g in pf.h.grids: g["Density"]

GPUARR = {}

with time_func("Copy to GPU"):
    for g in pf.h.grids:
        GPUARR[g.id] = cl_array.to_device(command_queue, g["Density"])
        cl.enqueue_barrier(command_queue)

with time_func("Sum on GPU1"):
    val_gpu1 = 0.0
    for g in pf.h.grids:
        val_gpu1 += cl_array.sum(GPUARR[g.id]).get()
        cl.enqueue_barrier(command_queue)

with time_func("Sum on GPU2"):
    val_gpu2 = 0.0
    for g in pf.h.grids:
        val_gpu2 += cl_array.sum(GPUARR[g.id]).get()
        cl.enqueue_barrier(command_queue)

#val_gpu = val_gpu.get()[0]
#print "On GPU calculated a sum of: %0.9e" % (val_gpu)

with time_func("De-allocate on GPU"):
    for g in pf.h.grids: del GPUARR[g.id]
    cl.enqueue_barrier(command_queue)

with time_func("Sum with NumPy"):
    val_cpu = sum(g["Density"].sum() * (g.RightEdge - g.LeftEdge).prod()
                  for g in pf.h.grids)