path: root/gpu.py

import numpy as np
import numpy.ma as ma
from copy import copy
from itertools import izip
import os
import sys

import pytools
import pycuda.tools
from pycuda.compiler import SourceModule
from pycuda import characterize
import pycuda.driver as cuda
from pycuda import gpuarray as ga

import chroma.src
from chroma.tools import timeit
from chroma.geometry import standard_wavelengths
from chroma.color import map_to_color
from chroma import event

cuda.init()

# standard nvcc options
cuda_options = ('--use_fast_math',)#, '--ptxas-options=-v']

@pycuda.tools.context_dependent_memoize
def get_cu_module(name, options=None, include_source_directory=True):
    """Returns a pycuda.compiler.SourceModule object from a CUDA source file
    located in the chroma src directory at src/[name].cu."""
    if options is None:
        options = []
    elif isinstance(options, tuple):
        options = list(options)
    else:
        raise TypeError('`options` must be a tuple.')

    srcdir = os.path.dirname(os.path.abspath(chroma.src.__file__))

    if include_source_directory:
        options += ['-I' + srcdir]

    with open('%s/%s' % (srcdir, name)) as f:
        source = f.read()

    return pycuda.compiler.SourceModule(source, options=options,
                                        no_extern_c=True)

def get_cu_source(name):
    srcdir = os.path.dirname(os.path.abspath(chroma.src.__file__))
    with open('%s/%s' % (srcdir, name)) as f:
        source = f.read()
    return source

class GPUFuncs(object):
    """Simple container class for GPU functions as attributes."""
    def __init__(self, module):
        self.module = module
        self.funcs = {}

    def __getattr__(self, name):
        try:
            return self.funcs[name]
        except KeyError:
            f = self.module.get_function(name)
            self.funcs[name] = f
            return f

init_rng_src = """
#include <curand_kernel.h>

extern "C"
{

__global__ void init_rng(int nthreads, curandState *s, unsigned long long seed, unsigned long long offset)
{
	int id = blockIdx.x*blockDim.x + threadIdx.x;

	if (id >= nthreads)
		return;

	curand_init(seed, id, offset, &s[id]);
}

} // extern "C"
"""

def get_rng_states(size, seed=1):
    "Return `size` number of CUDA random number generator states."
    rng_states = cuda.mem_alloc(size*characterize.sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))

    module = pycuda.compiler.SourceModule(init_rng_src, no_extern_c=True)
    init_rng = module.get_function('init_rng')

    init_rng(np.int32(size), rng_states, np.uint64(seed), np.uint64(0), block=(64,1,1), grid=(size//64+1,1))

    return rng_states

def to_float3(arr):
    "Returns an pycuda.gpuarray.vec.float3 array from an (N,3) array."
    if not arr.flags['C_CONTIGUOUS']:
        arr = np.asarray(arr, order='c')
    return arr.astype(np.float32).view(ga.vec.float3)[:,0]

def to_uint3(arr):
    "Returns a pycuda.gpuarray.vec.uint3 array from an (N,3) array."
    if not arr.flags['C_CONTIGUOUS']:
        arr = np.asarray(arr, order='c')
    return arr.astype(np.uint32).view(ga.vec.uint3)[:,0]

def chunk_iterator(nelements, nthreads_per_block=64, max_blocks=1024):
    """Iterator that yields tuples with the values requried to process
    a long array in multiple kernel passes on the GPU.

    Each yielded value is of the form,
        (first_index, elements_this_iteration, nblocks_this_iteration)

    Example:
        >>> list(chunk_iterator(300, 32, 2))
        [(0, 64, 2), (64, 64, 2), (128, 64, 2), (192, 64, 2), (256, 9, 1)]
    """
    first = 0
    while first < nelements:
        elements_left = nelements - first
        blocks = int(elements_left // nthreads_per_block)
        if elements_left % nthreads_per_block != 0:
            blocks += 1 # Round up only if needed
        blocks = min(max_blocks, blocks)
        elements_this_round = min(elements_left, blocks * nthreads_per_block)

        yield (first, elements_this_round, blocks)
        first += elements_this_round

class GPUPhotons(object):
    def __init__(self, photons):
        self.pos = ga.to_gpu(to_float3(photons.pos))
        self.dir = ga.to_gpu(to_float3(photons.dir))
        self.pol = ga.to_gpu(to_float3(photons.pol))
        self.wavelengths = ga.to_gpu(photons.wavelengths.astype(np.float32))
        self.t = ga.to_gpu(photons.t.astype(np.float32))
        self.last_hit_triangles = ga.to_gpu(photons.last_hit_triangles.astype(np.int32))
        self.flags = ga.to_gpu(photons.flags.astype(np.uint32))

        #cuda_options = ('--use_fast_math', '-w')#, '--ptxas-options=-v']

        module = get_cu_module('propagate.cu', options=cuda_options)
        self.gpu_funcs = GPUFuncs(module)

    def get(self):
        pos = self.pos.get().view(np.float32).reshape((len(self.pos),3))
        dir = self.dir.get().view(np.float32).reshape((len(self.dir),3))
        pol = self.pol.get().view(np.float32).reshape((len(self.pol),3))
        wavelengths = self.wavelengths.get()
        t = self.t.get()
        last_hit_triangles = self.last_hit_triangles.get()
        flags = self.flags.get()
        return event.Photons(pos, dir, pol, wavelengths, t, last_hit_triangles, flags)

    def propagate(self, gpu_geometry, rng_states, nthreads_per_block=64,
                  max_blocks=1024, max_steps=10):
        """Propagate photons on GPU to termination or max_steps, whichever
        comes first.

        May be called repeatedly without reloading photon information if
        single-stepping through photon history.

        ..warning::
            `rng_states` must have at least `nthreads_per_block`*`max_blocks`
            number of curandStates.
        """
        nphotons = self.pos.size
        step = 0
        input_queue = np.zeros(shape=nphotons+1, dtype=np.uint32)
        input_queue[1:] = np.arange(nphotons, dtype=np.uint32)
        input_queue_gpu = ga.to_gpu(input_queue)
        output_queue = np.zeros(shape=nphotons+1, dtype=np.uint32)
        output_queue[0] = 1
        output_queue_gpu = ga.to_gpu(output_queue)

        while step < max_steps:
            # Just finish the rest of the steps if the # of photons is low
            if nphotons < nthreads_per_block * 16 * 8:
                nsteps = max_steps - step
            else:
                nsteps = 1

            for first_photon, photons_this_round, blocks in \
                    chunk_iterator(nphotons, nthreads_per_block, max_blocks):
                self.gpu_funcs.propagate(np.int32(first_photon), np.int32(photons_this_round), input_queue_gpu[1:], output_queue_gpu, rng_states, self.pos, self.dir, self.wavelengths, self.pol, self.t, self.flags, self.last_hit_triangles, np.int32(nsteps), gpu_geometry.gpudata, block=(nthreads_per_block,1,1), grid=(blocks, 1))

            step += nsteps

            if step < max_steps:
                temp = input_queue_gpu
                input_queue_gpu = output_queue_gpu
                output_queue_gpu = temp
                output_queue_gpu[:1].set(np.uint32(1))
                nphotons = input_queue_gpu[:1].get()[0] - 1

        if ga.max(self.flags).get() & (1 << 31):
            print >>sys.stderr, "WARNING: ABORTED PHOTONS"

class GPUChannels(object):
    def __init__(self, t, q, flags):
        self.t = t
        self.q = q
        self.flags = flags

    def get(self):
        t = self.t.get()
        q = self.q.get().astype(np.float32)

        # For now, assume all channels with small
        # enough hit time were hit.
        return event.Channels(t<1e8, t, q, self.flags.get())

class GPURays(object):
    """The GPURays class holds arrays of ray positions and directions
    on the GPU that are used to render a geometry."""
    def __init__(self, pos, dir, max_alpha_depth=10, nblocks=64):
        self.pos = ga.to_gpu(to_float3(pos))
        self.dir = ga.to_gpu(to_float3(dir))

        self.max_alpha_depth = max_alpha_depth

        self.nblocks = nblocks

        transform_module = get_cu_module('transform.cu', options=cuda_options)
        self.transform_funcs = GPUFuncs(transform_module)

        render_module = get_cu_module('render.cu', options=cuda_options)
        self.render_funcs = GPUFuncs(render_module)

        self.dx = ga.empty(max_alpha_depth*self.pos.size, dtype=np.float32)
        self.color = ga.empty(self.dx.size, dtype=ga.vec.float4)
        self.dxlen = ga.zeros(self.pos.size, dtype=np.uint32)

    def rotate(self, phi, n):
        "Rotate by an angle phi around the axis `n`."
        self.transform_funcs.rotate(np.int32(self.pos.size), self.pos, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.pos.size//self.nblocks+1,1))
        self.transform_funcs.rotate(np.int32(self.dir.size), self.dir, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.dir.size//self.nblocks+1,1))

    def rotate_around_point(self, phi, n, point):
        """"Rotate by an angle phi around the axis `n` passing through
        the point `point`."""
        self.transform_funcs.rotate_around_point(np.int32(self.pos.size), self.pos, np.float32(phi), ga.vec.make_float3(*n), ga.vec.make_float3(*point), block=(self.nblocks,1,1), grid=(self.pos.size//self.nblocks+1,1))
        self.transform_funcs.rotate(np.int32(self.dir.size), self.dir, np.float32(phi), ga.vec.make_float3(*n), block=(self.nblocks,1,1), grid=(self.dir.size//self.nblocks+1,1))

    def translate(self, v):
        "Translate the ray positions by the vector `v`."
        self.transform_funcs.translate(np.int32(self.pos.size), self.pos, ga.vec.make_float3(*v), block=(self.nblocks,1,1), grid=(self.pos.size//self.nblocks+1,1))

    def render(self, gpu_geometry, pixels, alpha_depth=10,
               keep_last_render=False):
        """Render `gpu_geometry` and fill the GPU array `pixels` with pixel
        colors."""
        if not keep_last_render:
            self.dxlen.fill(0)

        if alpha_depth > self.max_alpha_depth:
            raise Exception('alpha_depth > max_alpha_depth')

        if not isinstance(pixels, ga.GPUArray):
            raise TypeError('`pixels` must be a %s instance.' % ga.GPUArray)

        if pixels.size != self.pos.size:
            raise ValueError('`pixels`.size != number of rays')

        self.render_funcs.render(np.int32(self.pos.size), self.pos, self.dir, gpu_geometry.gpudata, np.uint32(alpha_depth), pixels, self.dx, self.dxlen, self.color, block=(self.nblocks,1,1), grid=(self.pos.size//self.nblocks+1,1))

    def snapshot(self, gpu_geometry, alpha_depth=10):
        "Render `gpu_geometry` and return a numpy array of pixel colors."
        pixels = ga.empty(self.pos.size, dtype=np.uint32)
        self.render(gpu_geometry, pixels, alpha_depth)
        return pixels.get()

def make_gpu_struct(size, members):
    struct = cuda.mem_alloc(size)

    i = 0
    for member in members:
        if isinstance(member, ga.GPUArray):
            member = member.gpudata

        if isinstance(member, cuda.DeviceAllocation):
            if i % 8:
                raise Exception('cannot align 64-bit pointer. '
                                'arrange struct member variables in order of '
                                'decreasing size.')

            cuda.memcpy_htod(int(struct)+i, np.intp(int(member)))
            i += 8
        elif np.isscalar(member):
            cuda.memcpy_htod(int(struct)+i, member)
            i += member.nbytes
        else:
            raise TypeError('expected a GPU device pointer or scalar type.')

    return struct

def format_size(size):
    if size < 1e3:
        return '%.1f%s' % (size, ' ')
    elif size < 1e6:
        return '%.1f%s' % (size/1e3, 'K')
    elif size < 1e9:
        return '%.1f%s' % (size/1e6, 'M')
    else:
        return '%.1f%s' % (size/1e9, 'G')

def format_array(name, array):
    return '%-15s %6s %6s' % \
        (name, format_size(len(array)), format_size(array.nbytes))

class GPUGeometry(object):
    def __init__(self, geometry, wavelengths=None, print_usage=False):
        if wavelengths is None:
            wavelengths = standard_wavelengths

        try:
            wavelength_step = np.unique(np.diff(wavelengths)).item()
        except ValueError:
            raise ValueError('wavelengths must be equally spaced apart.')

        geometry_source = get_cu_source('geometry.h')
        material_struct_size = characterize.sizeof('Material', geometry_source)
        surface_struct_size = characterize.sizeof('Surface', geometry_source)
        geometry_struct_size = characterize.sizeof('Geometry', geometry_source)

        self.material_data = []
        self.material_ptrs = []

        def interp_material_property(wavelengths, property):
            # note that it is essential that the material properties be
            # interpolated linearly. this fact is used in the propagation
            # code to guarantee that probabilities still sum to one.
            return np.interp(wavelengths, property[:,0], property[:,1]).astype(np.float32)

        for i in range(len(geometry.unique_materials)):
            material = geometry.unique_materials[i]

            if material is None:
                raise Exception('one or more triangles is missing a material.')

            refractive_index = interp_material_property(wavelengths, material.refractive_index)
            refractive_index_gpu = ga.to_gpu(refractive_index)
            absorption_length = interp_material_property(wavelengths, material.absorption_length)
            absorption_length_gpu = ga.to_gpu(absorption_length)
            scattering_length = interp_material_property(wavelengths, material.scattering_length)
            scattering_length_gpu = ga.to_gpu(scattering_length)

            self.material_data.append(refractive_index_gpu)
            self.material_data.append(absorption_length_gpu)
            self.material_data.append(scattering_length_gpu)

            material_gpu = \
                make_gpu_struct(material_struct_size,
                                [refractive_index_gpu, absorption_length_gpu,
                                 scattering_length_gpu,
                                 np.uint32(len(wavelengths)),
                                 np.float32(wavelength_step),
                                 np.float32(wavelengths[0])])

            self.material_ptrs.append(material_gpu)

        self.material_pointer_array = \
            make_gpu_struct(8*len(self.material_ptrs), self.material_ptrs)

        self.surface_data = []
        self.surface_ptrs = []

        for i in range(len(geometry.unique_surfaces)):
            surface = geometry.unique_surfaces[i]

            if surface is None:
                # need something to copy to the surface array struct
                # that is the same size as a 64-bit pointer.
                # this pointer will never be used by the simulation.
                self.surface_ptrs.append(np.uint64(0))
                continue

            detect = interp_material_property(wavelengths, surface.detect)
            detect_gpu = ga.to_gpu(detect)
            absorb = interp_material_property(wavelengths, surface.absorb)
            absorb_gpu = ga.to_gpu(absorb)
            reflect_diffuse = interp_material_property(wavelengths, surface.reflect_diffuse)
            reflect_diffuse_gpu = ga.to_gpu(reflect_diffuse)
            reflect_specular = interp_material_property(wavelengths, surface.reflect_specular)
            reflect_specular_gpu = ga.to_gpu(reflect_specular)

            self.surface_data.append(detect_gpu)
            self.surface_data.append(absorb_gpu)
            self.surface_data.append(reflect_diffuse_gpu)
            self.surface_data.append(reflect_specular_gpu)

            surface_gpu = \
                make_gpu_struct(surface_struct_size,
                                [detect_gpu, absorb_gpu,
                                 reflect_diffuse_gpu,
                                 reflect_specular_gpu,
                                 np.uint32(len(wavelengths)),
                                 np.float32(wavelength_step),
                                 np.float32(wavelengths[0])])

            self.surface_ptrs.append(surface_gpu)

        self.surface_pointer_array = \
            make_gpu_struct(8*len(self.surface_ptrs), self.surface_ptrs)

        self.vertices = ga.to_gpu(to_float3(geometry.mesh.vertices))
        self.triangles = ga.to_gpu(to_uint3(geometry.mesh.triangles))

        material_codes = (((geometry.material1_index & 0xff) << 24) |
                          ((geometry.material2_index & 0xff) << 16) |
                          ((geometry.surface_index & 0xff) << 8)).astype(np.uint32)

        self.material_codes = ga.to_gpu(material_codes)

        self.lower_bounds = ga.to_gpu(to_float3(geometry.lower_bounds))
        self.upper_bounds = ga.to_gpu(to_float3(geometry.upper_bounds))
        self.colors = ga.to_gpu(geometry.colors.astype(np.uint32))
        self.node_map = ga.to_gpu(geometry.node_map.astype(np.uint32))
        self.node_map_end = ga.to_gpu(geometry.node_map_end.astype(np.uint32))
        self.solid_id_map = ga.to_gpu(geometry.solid_id.astype(np.uint32))

        self.gpudata = make_gpu_struct(geometry_struct_size,
                                       [self.vertices, self.triangles,
                                        self.material_codes,
                                        self.colors, self.lower_bounds,
                                        self.upper_bounds, self.node_map,
                                        self.node_map_end,
                                        self.material_pointer_array,
                                        self.surface_pointer_array,
                                        np.uint32(geometry.start_node),
                                        np.uint32(geometry.first_node)])

        self.geometry = geometry

        if print_usage:
            self.print_device_usage()

    def print_device_usage(self):
        print 'device usage:'
        print '-'*10
        print format_array('vertices', self.vertices)
        print format_array('triangles', self.triangles)
        print format_array('lower_bounds', self.lower_bounds)
        print format_array('upper_bounds', self.upper_bounds)
        print format_array('node_map', self.node_map)
        print format_array('node_map_end', self.node_map_end)
        print '%-15s %6s %6s' % ('total', '', format_size(self.vertices.nbytes + self.triangles.nbytes + self.lower_bounds.nbytes + self.upper_bounds.nbytes + self.node_map.nbytes + self.node_map_end.nbytes))
        print '-'*10
        free, total = cuda.mem_get_info()
        print '%-15s %6s %6s' % ('device total', '', format_size(total))
        print '%-15s %6s %6s' % ('device used', '', format_size(total-free))
        print '%-15s %6s %6s' % ('device free', '', format_size(free))
        print

    def reset_colors(self):
        self.colors.set_async(self.geometry.colors.astype(np.uint32))

    def color_solids(self, solid_hit, colors, nblocks_per_thread=64,
                     max_blocks=1024):
        solid_hit_gpu = ga.to_gpu(np.array(solid_hit, dtype=np.bool))
        solid_colors_gpu = ga.to_gpu(np.array(colors, dtype=np.uint32))

        module = get_cu_module('mesh.h', options=cuda_options)
        color_solids = module.get_function('color_solids')

        for first_triangle, triangles_this_round, blocks in \
                chunk_iterator(self.triangles.size, nblocks_per_thread,
                               max_blocks):
            color_solids(np.int32(first_triangle),
                         np.int32(triangles_this_round), self.solid_id_map,
                         solid_hit_gpu, solid_colors_gpu, self.gpudata,
                         block=(nblocks_per_thread,1,1), 
                         grid=(blocks,1))

class GPUDaq(object):
    def __init__(self, gpu_geometry, max_pmt_id, pmt_rms=1.2e-9):
        self.earliest_time_gpu = ga.empty(max_pmt_id+1, dtype=np.float32)
        self.earliest_time_int_gpu = ga.empty(max_pmt_id+1, dtype=np.uint32)
        self.channel_history_gpu = ga.zeros_like(self.earliest_time_int_gpu)
        self.channel_q_gpu = ga.zeros_like(self.earliest_time_int_gpu)
        self.daq_pmt_rms = pmt_rms
        self.solid_id_map_gpu = gpu_geometry.solid_id_map

        self.module = get_cu_module('daq.cu', options=cuda_options, 
                                    include_source_directory=False)
        self.gpu_funcs = GPUFuncs(self.module)

    def acquire(self, gpuphotons, rng_states, nthreads_per_block=64, max_blocks=1024):
        self.gpu_funcs.reset_earliest_time_int(np.float32(1e9), np.int32(len(self.earliest_time_int_gpu)), self.earliest_time_int_gpu, block=(nthreads_per_block,1,1), grid=(len(self.earliest_time_int_gpu)//nthreads_per_block+1,1))
        self.channel_q_gpu.fill(0)
        self.channel_history_gpu.fill(0)

        n = len(gpuphotons.pos)

        for first_photon, photons_this_round, blocks in \
                chunk_iterator(n, nthreads_per_block, max_blocks):
            self.gpu_funcs.run_daq(rng_states, np.uint32(0x1 << 2), np.float32(self.daq_pmt_rms), np.int32(first_photon), np.int32(photons_this_round), gpuphotons.t, gpuphotons.flags, gpuphotons.last_hit_triangles, self.solid_id_map_gpu, np.int32(len(self.earliest_time_int_gpu)), self.earliest_time_int_gpu, self.channel_q_gpu, self.channel_history_gpu, block=(nthreads_per_block,1,1), grid=(blocks,1))

        self.gpu_funcs.convert_sortable_int_to_float(np.int32(len(self.earliest_time_int_gpu)), self.earliest_time_int_gpu, self.earliest_time_gpu, block=(nthreads_per_block,1,1), grid=(len(self.earliest_time_int_gpu)//nthreads_per_block+1,1))

        return GPUChannels(self.earliest_time_gpu, self.channel_q_gpu, self.channel_history_gpu)

class GPUPDF(object):
    def __init__(self):
        self.module = get_cu_module('daq.cu', options=cuda_options,
                                    include_source_directory=False)
        self.gpu_funcs = GPUFuncs(self.module)

    def setup_pdf(self, max_pmt_id, tbins, trange, qbins, qrange):
        """Setup GPU arrays to hold PDF information.

           max_pmt_id: int, largest PMT id #
           tbins: number of time bins
           trange: tuple of (min, max) time in PDF
           qbins: number of charge bins
           qrange: tuple of (min, max) charge in PDF
        """
        self.events_in_histogram = 0
        self.hitcount_gpu = ga.zeros(max_pmt_id+1, dtype=np.uint32)
        self.pdf_gpu = ga.zeros(shape=(max_pmt_id+1, tbins, qbins), 
                                      dtype=np.uint32)
        self.tbins = tbins
        self.trange = trange
        self.qbins = qbins
        self.qrange = qrange

    def clear_pdf(self):
        """Rezero the PDF counters."""
        self.hitcount_gpu.fill(0)
        self.pdf_gpu.fill(0)

    def add_hits_to_pdf(self, gpuchannels, nthreads_per_block=64):
        self.gpu_funcs.bin_hits(np.int32(len(self.hitcount_gpu)),
                                gpuchannels.q,
                                gpuchannels.t,
                                self.hitcount_gpu,
                                np.int32(self.tbins),
                                np.float32(self.trange[0]),
                                np.float32(self.trange[1]),
                                np.int32(self.qbins),
                                np.float32(self.qrange[0]),
                                np.float32(self.qrange[1]),
                                self.pdf_gpu,
                                block=(nthreads_per_block,1,1), 
                                grid=(len(gpuchannels.t)//nthreads_per_block+1,1))


        self.events_in_histogram += 1

    def get_pdfs(self):
        """Returns the 1D hitcount array and the 3D [channel, time, charge]
        histogram."""
        return self.hitcount_gpu.get(), self.pdf_gpu.get()

    def setup_pdf_eval(self, event_hit, event_time, event_charge, min_twidth,
                       trange, min_qwidth, qrange, min_bin_content=10,
                       time_only=True):
        """Setup GPU arrays to compute PDF values for the given event.
        The pdf_eval calculation allows the PDF to be evaluated at a
        single point for each channel as the Monte Carlo is run.  The
        effective bin size will be as small as (`min_twidth`,
        `min_qwidth`) around the point of interest, but will be large
        enough to ensure that `min_bin_content` Monte Carlo events
        fall into the bin.

            event_hit: ndarray
              Hit or not-hit status for each channel in the detector.
            event_time: ndarray
              Hit time for each channel in the detector.  If channel 
              not hit, the time will be ignored.
            event_charge: ndarray
              Integrated charge for each channel in the detector.
              If channel not hit, the charge will be ignored.

            min_twidth: float
              Minimum bin size in the time dimension
            trange: (float, float)
              Range of time dimension in PDF
            min_qwidth: float
              Minimum bin size in charge dimension
            qrange: (float, float)
              Range of charge dimension in PDF
            min_bin_content: int
              The bin will be expanded to include at least this many events
            time_only: bool
              If True, only the time observable will be used in the PDF.
        """
        self.event_hit_gpu = ga.to_gpu(event_hit.astype(np.uint32))
        self.event_time_gpu = ga.to_gpu(event_time.astype(np.float32))
        self.event_charge_gpu = ga.to_gpu(event_charge.astype(np.float32))

        self.eval_hitcount_gpu = ga.zeros(len(event_hit), dtype=np.uint32)
        self.eval_bincount_gpu = ga.zeros(len(event_hit), dtype=np.uint32)
        self.nearest_mc_gpu = ga.empty(shape=len(event_hit) * min_bin_content, 
                                             dtype=np.float32)
        self.nearest_mc_gpu.fill(1e9)

        self.min_twidth = min_twidth
        self.trange = trange
        self.min_qwidth = min_qwidth
        self.qrange = qrange
        self.min_bin_content = min_bin_content
        self.time_only = time_only

    def clear_pdf_eval(self):
        "Reset PDF evaluation counters to start accumulating new Monte Carlo."
        self.eval_hitcount_gpu.fill(0)
        self.eval_bincount_gpu.fill(0)
        self.nearest_mc_gpu.fill(1e9)

    def accumulate_pdf_eval(self, gpuchannels, nthreads_per_block=64):
        "Add the most recent results of run_daq() to the PDF evaluation."
        self.gpu_funcs.accumulate_pdf_eval(np.int32(self.time_only),
                                           np.int32(len(self.event_hit_gpu)),
                                           self.event_hit_gpu,
                                           self.event_time_gpu,
                                           self.event_charge_gpu,
                                           gpuchannels.t,
                                           gpuchannels.q,
                                           self.eval_hitcount_gpu,
                                           self.eval_bincount_gpu,
                                           np.float32(self.min_twidth),
                                           np.float32(self.trange[0]),
                                           np.float32(self.trange[1]),
                                           np.float32(self.min_qwidth),
                                           np.float32(self.qrange[0]),
                                           np.float32(self.qrange[1]),
                                           self.nearest_mc_gpu,
                                           np.int32(self.min_bin_content),
                                           block=(nthreads_per_block,1,1), 
                                           grid=(len(gpuchannels.t)//nthreads_per_block+1,1))

    def get_pdf_eval(self):
        evhit = self.event_hit_gpu.get().astype(bool)
        hitcount = self.eval_hitcount_gpu.get()
        bincount = self.eval_bincount_gpu.get()

        pdf_value = np.zeros(len(hitcount), dtype=float)
        pdf_frac_uncert = np.zeros_like(pdf_value)

        # PDF value for high stats bins
        high_stats = bincount >= self.min_bin_content
        if high_stats.any():
            if self.time_only:
                pdf_value[high_stats] = bincount[high_stats].astype(float) / hitcount[high_stats] / self.min_twidth
            else:
                assert Exception('Unimplemented 2D (time,charge) mode!')

            pdf_frac_uncert[high_stats] = 1.0/np.sqrt(bincount[high_stats])

        # PDF value for low stats bins
        low_stats = ~high_stats & (hitcount > 0) & evhit

        nearest_mc = self.nearest_mc_gpu.get().reshape((len(hitcount), self.min_bin_content))

        # Deal with the case where we did not even get min_bin_content events
        # in the PDF but also clamp the lower range to ensure we don't index
        # by a negative number in 2 lines
        last_valid_entry = np.maximum(0, (nearest_mc < 1e9).astype(int).sum(axis=1) - 1)
        distance = nearest_mc[np.arange(len(last_valid_entry)),last_valid_entry]

        if low_stats.any():
            if self.time_only:
                pdf_value[low_stats] = (last_valid_entry[low_stats] + 1).astype(float) / hitcount[low_stats] / distance[low_stats] / 2.0
            else:
                assert Exception('Unimplemented 2D (time,charge) mode!')

            pdf_frac_uncert[low_stats] = 1.0/np.sqrt(last_valid_entry[low_stats] + 1)

        # PDFs with no stats got zero by default during array creation
        
        print 'high_stats:', high_stats.sum(), 'low_stats', low_stats.sum()

        return hitcount, pdf_value, pdf_value * pdf_frac_uncert

def create_context(device_id=None):
    """Initialize and return a GPU context on the specified device.
    If device_id is None, the default device is used."""
    try:
        cuda.mem_get_info()
    except cuda.LogicError:
        # initialize cuda
        cuda.init()

    if device_id is None:
        context = pycuda.tools.make_default_context()
    else:
        device = cuda.Device(device_id)
        context = device.make_context()

    context.set_cache_config(cuda.func_cache.PREFER_L1)

    return context

class GPU(object):
    def __init__(self, device_id=None):
        """Initialize a GPU context on the specified device.
        If device_id is None, the default device is used."""

        if device_id is None:
            self.context = pycuda.tools.make_default_context()
        else:
            device = cuda.Device(device_id)
            self.context = device.make_context()

        print 'device %s' % self.context.get_device().name()

        self.print_mem_info()

        self.context.set_cache_config(cuda.func_cache.PREFER_L1)

    def print_mem_info(self):
        free, total = cuda.mem_get_info()

        print '%.1f %% of device memory is free.' % ((free/float(total))*100)

    def __del__(self):
        self.context.pop()