diff options
| -rw-r--r-- | chroma/gpu/geometry.py | 37 |
1 files changed, 22 insertions, 15 deletions
diff --git a/chroma/gpu/geometry.py b/chroma/gpu/geometry.py index a223006..081563f 100644 --- a/chroma/gpu/geometry.py +++ b/chroma/gpu/geometry.py @@ -7,6 +7,7 @@ from chroma.geometry import standard_wavelengths from chroma.gpu.tools import get_cu_module, get_cu_source, cuda_options, \ chunk_iterator, format_array, format_size, to_uint3, to_float3, \ make_gpu_struct +from chroma.log import logger class GPUGeometry(object): def __init__(self, geometry, wavelengths=None, print_usage=False): @@ -134,23 +135,29 @@ class GPUGeometry(object): if print_usage: self.print_device_usage() + logger.info(self.device_usage_str()) + + def device_usage_str(self): + '''Returns a formatted string displaying the memory usage.''' + s = 'device usage:\n' + s += '-'*10 + '\n' + s += format_array('vertices', self.vertices) + '\n' + s += format_array('triangles', self.triangles) + '\n' + s += format_array('lower_bounds', self.lower_bounds) + '\n' + s += format_array('upper_bounds', self.upper_bounds) + '\n' + s += format_array('node_map', self.node_map) + '\n' + s += format_array('node_map_end', self.node_map_end) + '\n' + s += '%-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)) + '\n' + s += '-'*10 + '\n' + free, total = cuda.mem_get_info() + s += '%-15s %6s %6s' % ('device total', '', format_size(total)) + '\n' + s += '%-15s %6s %6s' % ('device used', '', format_size(total-free)) + '\n' + s += '%-15s %6s %6s' % ('device free', '', format_size(free)) + '\n' + return s 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 + print self.device_usage_str() + print def reset_colors(self): self.colors.set_async(self.geometry.colors.astype(np.uint32)) |