path: root/gpu.py

import numpy as np
import numpy.ma as ma
from pycuda.tools import make_default_context
from pycuda.compiler import SourceModule
from pycuda.characterize import sizeof
import pycuda.driver as cuda
from pycuda import gpuarray
from copy import copy
from itertools import izip
from tools import timeit

import src
from geometry import standard_wavelengths
from color import map_to_color

cuda.init()

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

       Each yielded value is (first_index, elements_this_iteration, nblocks_this_iteration)

    >>> 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 = 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
    

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 CUDAFuncs(object):
    '''simple container class for GPU functions as attributes'''
    def __init__(self, cuda_module, func_names):
        '''Extract __global__ functions listed in func_names from
        the PyCUDA module object.  The functions are assigned
        to attributes of this object with the same name.'''
        for name in func_names:
            setattr(self, name, cuda_module.get_function(name))

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

        if device_id is None:
            self.context = make_default_context()
        else:
            print 'device_id = %s' % device_id
            device = cuda.Device(device_id)
            self.context = device.make_context()
        print 'device %s' % self.context.get_device().name()

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

        cuda_options = ['-I' + src.dir, '--use_fast_math', '--ptxas-options=-v']

        self.module = SourceModule(src.kernel, options=cuda_options, no_extern_c=True)
        self.geo_funcs = CUDAFuncs(self.module, ['set_wavelength_range', 'set_material', 'set_surface', 'set_global_mesh_variables', 'color_solids'])

        self.prop_funcs = CUDAFuncs(self.module, ['init_rng', 'propagate'])
        self.nthread_per_block = 64
        self.max_blocks = 1024
        self.daq_module = SourceModule(src.daq, options=cuda_options, no_extern_c=True)
        self.daq_funcs = CUDAFuncs(self.daq_module, 
                              ['reset_earliest_time_int', 'run_daq',
                               'convert_sortable_int_to_float'])

    def print_device_usage(self):
        print 'device usage:'
        print format_array('vertices', self.vertices_gpu)
        print format_array('triangles', self.triangles_gpu)
        print format_array('lower_bounds', self.lower_bounds_gpu)
        print format_array('upper_bounds', self.upper_bounds_gpu)
        print format_array('node_map', self.node_map_gpu)
        print format_array('node_map_end', self.node_map_end_gpu)
        print '%-15s %6s %6s' % ('total', '', format_size(self.vertices_gpu.nbytes + self.triangles_gpu.nbytes + self.lower_bounds_gpu.nbytes + self.upper_bounds_gpu.nbytes + self.node_map_gpu.nbytes + self.node_map_end_gpu.nbytes))

    def load_geometry(self, geometry):
        if not hasattr(geometry, 'mesh'):
            print 'WARNING: building geometry with 8-bits'
            geometry.build(bits=8)

        self.geo_funcs.set_wavelength_range(np.float32(standard_wavelengths[0]), np.float32(standard_wavelengths[-1]), np.float32(standard_wavelengths[1]-standard_wavelengths[0]), np.uint32(standard_wavelengths.size), block=(1,1,1), grid=(1,1))

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

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

            material.refractive_index_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, material.refractive_index[:,0], material.refractive_index[:,1]).astype(np.float32))
            material.absorption_length_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, material.absorption_length[:,0], material.absorption_length[:,1]).astype(np.float32))
            material.scattering_length_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, material.scattering_length[:,0], material.scattering_length[:,1]).astype(np.float32))

            self.geo_funcs.set_material(np.int32(i), material.refractive_index_gpu, material.absorption_length_gpu, material.scattering_length_gpu, block=(1,1,1), grid=(1,1))

            self.materials.append(material)

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

            if surface is None:
                continue

            surface.detect_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.detect[:,0], surface.detect[:,1]).astype(np.float32))
            surface.absorb_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.absorb[:,0], surface.absorb[:,1]).astype(np.float32))
            surface.reflect_diffuse_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.reflect_diffuse[:,0], surface.reflect_diffuse[:,1]).astype(np.float32))
            surface.reflect_specular_gpu = gpuarray.to_gpu(np.interp(standard_wavelengths, surface.reflect_specular[:,0], surface.reflect_specular[:,1]).astype(np.float32))

            self.geo_funcs.set_surface(np.int32(i), surface.detect_gpu, surface.absorb_gpu, surface.reflect_diffuse_gpu, surface.reflect_specular_gpu, block=(1,1,1), grid=(1,1))

            self.surfaces.append(surface)

        self.vertices_gpu = gpuarray.to_gpu(geometry.mesh.vertices.astype(np.float32).view(gpuarray.vec.float3))

        triangles = \
            np.empty(len(geometry.mesh.triangles), dtype=gpuarray.vec.uint4)
        triangles['x'] = geometry.mesh.triangles[:,0]
        triangles['y'] = geometry.mesh.triangles[:,1]
        triangles['z'] = geometry.mesh.triangles[:,2]
        triangles['w'] = ((geometry.material1_index & 0xff) << 24) | ((geometry.material2_index & 0xff) << 16) | ((geometry.surface_index & 0xff) << 8)
        self.triangles_gpu = gpuarray.to_gpu(triangles)

        lower_bounds_float3 = np.empty(geometry.lower_bounds.shape[0], dtype=gpuarray.vec.float3)
        lower_bounds_float3['x'] = geometry.lower_bounds[:,0]
        lower_bounds_float3['y'] = geometry.lower_bounds[:,1]
        lower_bounds_float3['z'] = geometry.lower_bounds[:,2]
        self.lower_bounds_gpu = gpuarray.to_gpu(lower_bounds_float3)

        upper_bounds_float3 = np.empty(geometry.upper_bounds.shape[0], dtype=gpuarray.vec.float3)
        upper_bounds_float3['x'] = geometry.upper_bounds[:,0]
        upper_bounds_float3['y'] = geometry.upper_bounds[:,1]
        upper_bounds_float3['z'] = geometry.upper_bounds[:,2]
        self.upper_bounds_gpu = gpuarray.to_gpu(upper_bounds_float3)

        self.colors_gpu = gpuarray.to_gpu(geometry.colors.astype(np.uint32))
        self.node_map_gpu = gpuarray.to_gpu(geometry.node_map.astype(np.uint32))
        self.node_map_end_gpu = gpuarray.to_gpu(geometry.node_map_end.astype(np.uint32))
        self.solid_id_map_gpu = gpuarray.to_gpu(geometry.solid_id.astype(np.uint32))

        self.geo_funcs.set_global_mesh_variables(self.triangles_gpu, self.vertices_gpu, self.colors_gpu, np.uint32(geometry.node_map.size-1), np.uint32(geometry.first_node), self.lower_bounds_gpu, self.upper_bounds_gpu, block=(1,1,1), grid=(1,1))

        self.node_map_tex = self.module.get_texref('node_map')
        self.node_map_end_tex = self.module.get_texref('node_map_end')

        self.node_map_tex.set_address(self.node_map_gpu.gpudata, self.node_map_gpu.nbytes)
        self.node_map_end_tex.set_address(self.node_map_end_gpu.gpudata, self.node_map_end_gpu.nbytes)

        self.node_map_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)
        self.node_map_end_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1)

        self.geometry = geometry

        print 'average of %f child nodes per node' % (np.mean(geometry.node_map_end[geometry.first_node:] - geometry.node_map[geometry.first_node:]))

        print '%i nodes with one child' % (np.count_nonzero((geometry.node_map_end[geometry.first_node:] - geometry.node_map[geometry.first_node:]) == 1))

        print '%i leaf nodes with one child' % (np.count_nonzero((geometry.node_map_end[:geometry.first_node] - geometry.node_map[:geometry.first_node]) == 1))


        self.print_device_usage()

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

    def color_solids(self, solid_ids, colors):
        solid_ids_gpu = gpuarray.to_gpu(np.array(solid_ids, dtype=np.int32))
        solid_colors_gpu = gpuarray.to_gpu(np.array(colors, dtype=np.uint32))

        self.geo_funcs.color_solids(np.int32(solid_ids_gpu.size), np.uint32(self.triangles_gpu.size), self.solid_id_map_gpu, solid_ids_gpu, solid_colors_gpu, block=(self.nthread_per_block,1,1), grid=(solid_ids_gpu.size//self.nthread_per_block+1,1))
        self.context.synchronize()

    def setup_propagate(self, seed=1):
        self.rng_states_gpu = cuda.mem_alloc(self.nthread_per_block * self.max_blocks
                                             * sizeof('curandStateXORWOW', '#include <curand_kernel.h>'))
        self.prop_funcs.init_rng(np.int32(self.max_blocks * self.nthread_per_block), self.rng_states_gpu, 
                                 np.uint64(seed), np.uint64(0), 
                                 block=(self.nthread_per_block,1,1), 
                                 grid=(self.max_blocks,1))
        #self.context.synchronize()

    def load_photons(self, pos, dir, pol, wavelength, t0, 
                     histories=None, last_hit_triangles=None):
        '''Load N photons onto the GPU.

           pos: numpy.array(shape=(N, 3)) of photon starting positions (meters)
           dir: numpy.array(shape=(N, 3)) of photon starting directions (unit vectors)
           pol: numpy.array(shape=(N, 3)) of photon polarization directions (unit vectors)
           wavelength: numpy.array(shape=N) of photon wavelengths (nm)
           t0: numpy.array(shape=N) of photon start times (s)
           
           Optional args will be loaded with defaults on GPU if not set:
           histories: Bitmask of interactions that have occurred over history of photon
           last_hit_triangles: The triangle ID number that the photon last interacted with,
                               if any.  -1 if no triangle was hit in the last step
        '''
        self.nphotons = len(pos)
        assert len(dir) == self.nphotons
        assert len(pol) == self.nphotons
        assert len(wavelength) == self.nphotons
        assert len(t0) == self.nphotons

        self.positions_gpu = gpuarray.to_gpu(pos.astype(np.float32).view(gpuarray.vec.float3))
        self.directions_gpu = gpuarray.to_gpu(dir.astype(np.float32).view(gpuarray.vec.float3))
        self.polarizations_gpu = gpuarray.to_gpu(pol.astype(np.float32).view(gpuarray.vec.float3))
        self.wavelengths_gpu = gpuarray.to_gpu(wavelength.astype(np.float32))
        self.times_gpu = gpuarray.to_gpu(t0.astype(np.float32))
        self.photon_offsets_gpu = gpuarray.arange(self.nphotons, dtype=np.int32)

        if histories is not None:
            self.histories_gpu = gpuarray.to_gpu(histories.astype(np.uint32))
        else:
            self.histories_gpu = gpuarray.zeros(self.nphotons, dtype=np.uint32)

        if last_hit_triangles is not None:
            self.last_hit_triangles_gpu = gpuarray.to_gpu(last_hit_triangles.astype(np.int32))
        else:
            self.last_hit_triangles_gpu = gpuarray.empty(self.nphotons, dtype=np.int32)
            self.last_hit_triangles_gpu.fill(-1)


    def propagate(self, 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.
        '''
        for first_photon, photons_this_round, blocks in chunk_iterator(self.nphotons, self.nthread_per_block, self.max_blocks):
            self.prop_funcs.propagate(np.int32(first_photon),
                                      np.int32(photons_this_round),
                                      self.photon_offsets_gpu,
                                      self.rng_states_gpu, self.positions_gpu, 
                                      self.directions_gpu, 
                                      self.wavelengths_gpu, self.polarizations_gpu, 
                                      self.times_gpu, self.histories_gpu, 
                                      self.last_hit_triangles_gpu, 
                                      np.int32(max_steps), 
                                      block=(self.nthread_per_block,1,1),
                                      grid=(blocks, 1))

        if 'profile' in __builtins__:
            self.context.synchronize()

    def get_photons(self):
        '''Returns a dictionary of current photon state information.

        Contents of dictionary have the same names as the parameters to load_photons().
        '''
        return { 'pos' : self.positions_gpu.get().view(np.float32).reshape(self.positions_gpu.size, 3),
                 'dir' : self.directions_gpu.get().view(np.float32).reshape(self.directions_gpu.size, 3),
                 'pol' : self.polarizations_gpu.get().view(np.float32).reshape(self.polarizations_gpu.size, 3),
                 'wavelength' : self.wavelengths_gpu.get(),
                 't0' : self.times_gpu.get(),
                 'histories' : self.histories_gpu.get(),
                 'last_hit_triangles' : self.last_hit_triangles_gpu.get()}
    
    def setup_daq(self, max_pmt_id, pmt_rms=1.2e-9):
        self.earliest_time_gpu = gpuarray.GPUArray(shape=(max_pmt_id+1,), dtype=np.float32)
        self.earliest_time_int_gpu = gpuarray.GPUArray(shape=self.earliest_time_gpu.shape, 
                                                       dtype=np.uint32)
        self.channel_history_gpu = gpuarray.zeros_like(self.earliest_time_int_gpu)
        self.channel_q_gpu = gpuarray.zeros_like(self.earliest_time_int_gpu)
        self.daq_pmt_rms = pmt_rms

    def run_daq(self):
        self.daq_funcs.reset_earliest_time_int(np.float32(1e9), 
                                               np.int32(len(self.earliest_time_int_gpu)), 
                                               self.earliest_time_int_gpu, 
                                               block=(self.nthread_per_block,1,1), 
                                               grid=(len(self.earliest_time_int_gpu)//self.nthread_per_block+1,1))
        self.channel_q_gpu.fill(0)
        self.channel_history_gpu.fill(0)

        #self.context.synchronize()
        for first_photon, photons_this_round, blocks in chunk_iterator(self.nphotons, self.nthread_per_block, self.max_blocks):
            self.daq_funcs.run_daq(self.rng_states_gpu, np.uint32(0x1 << 2), 
                                   np.float32(self.daq_pmt_rms), 
                                   np.int32(first_photon),
                                   np.int32(photons_this_round), 
                                   self.times_gpu, self.histories_gpu,
                                   self.last_hit_triangles_gpu, 
                                   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=(self.nthread_per_block,1,1), grid=(blocks,1))
        #self.context.synchronize()
        self.daq_funcs.convert_sortable_int_to_float(np.int32(len(self.earliest_time_int_gpu)), 
                                                     self.earliest_time_int_gpu, 
                                                     self.earliest_time_gpu, 
                                                     block=(self.nthread_per_block,1,1), 
                                                     grid=(len(self.earliest_time_int_gpu)//self.nthread_per_block+1,1))
        if 'profile' in __builtins__:
            self.context.synchronize()
        

    def get_hits(self):
        return { 't': self.earliest_time_gpu.get(), 
                 'q': self.channel_q_gpu.get().astype(np.float32),
                 'history': self.channel_history_gpu.get()}

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