Improve CUDA array header allocation

Author: sklam

numba/cuda/cudadrv/devicearray.py

@@ -7,7 +7,7 @@
 import warnings
 import math
 import numpy as np
-from .ndarray import ndarray_device_allocate_head, ndarray_populate_head
+from .ndarray import (ndarray_populate_head, ArrayHeaderManager)
 from . import driver as _driver
 from . import devices
 from numba import dummyarray
@@ -54,15 +54,23 @@ class DeviceNDArrayBase(object):
     def __init__(self, shape, strides, dtype, stream=0, writeback=None,
                  gpu_head=None, gpu_data=None):
         """
-        Arguments
-
-        shape: array shape.
-        strides: array strides.
-        dtype: data type as numpy.dtype.
-        stream: cuda stream.
-        writeback: Deprecated.
-        gpu_head: user provided device memory for the ndarray head structure
-        gpu_data: user provided device memory for the ndarray data buffer
+        Args
+        ----
+
+        shape
+            array shape.
+        strides
+            array strides.
+        dtype
+            data type as numpy.dtype.
+        stream
+            cuda stream.
+        writeback
+            Deprecated.
+        gpu_head
+            user provided device memory for the ndarray head structure
+        gpu_data
+            user provided device memory for the ndarray data buffer
         """
         if isinstance(shape, (int, long)):
             shape = (shape,)
@@ -86,27 +94,22 @@ def __init__(self, shape, strides, dtype, stream=0, writeback=None,
         else:
             self.alloc_size = _driver.device_memory_size(gpu_data)
 
+        self.gpu_mem = ArrayHeaderManager(devices.get_context())
+
         if gpu_head is None:
-            gpu_head = ndarray_device_allocate_head(self.ndim)
+            gpu_head = self.gpu_mem.allocate(self.ndim)
             ndarray_populate_head(gpu_head, gpu_data, self.shape,
                                   self.strides, stream=stream)
         self.gpu_head = gpu_head
         self.gpu_data = gpu_data
 
         self.__writeback = writeback # should deprecate the use of this
 
-        # define the array interface to work with numpy
-        #
-        # XXX: problem with data being accessed.
-        #      is NULL pointer alright?
-        #
-        #        self.__array_interface__ = {
-        #            'shape'     : self.shape,
-        #            'typestr'   : self.dtype.str,
-        #            'data'      : (0, True),
-        #            'version'   : 3,
-        #        }
-
+    def __del__(self):
+        try:
+            self.gpu_mem.free(self.gpu_head)
+        except:
+            pass
 
     @property
     def device_ctypes_pointer(self):
@@ -183,11 +186,11 @@ def split(self, section, stream=0):
             end = min(begin + section, self.size)
             shape = (end - begin,)
             gpu_data = self.gpu_data.view(begin * itemsize, end * itemsize)
-            gpu_head = ndarray_device_allocate_head(1)
-            ndarray_populate_head(gpu_head, gpu_data, shape, strides,
-                                  stream=stream)
+            # gpu_head = _allocate_head(1)
+            # ndarray_populate_head(gpu_head, gpu_data, shape, strides,
+            #                       stream=stream)
             yield DeviceNDArray(shape, strides, dtype=self.dtype, stream=stream,
-                                gpu_head=gpu_head, gpu_data=gpu_data)
+                                gpu_data=gpu_data)
 
     def as_cuda_arg(self):
         """Returns a device memory object that is used as the argument.
@@ -251,14 +254,21 @@ def __getitem__(self, item):
 
 class MappedNDArray(DeviceNDArrayBase, np.ndarray):
     def device_setup(self, gpu_data, stream=0):
-        gpu_head = ndarray_device_allocate_head(self.ndim)
+        self.gpu_mem = ArrayHeaderManager(devices.get_context())
 
+        gpu_head = self.gpu_mem.allocate(self.ndim)
         ndarray_populate_head(gpu_head, gpu_data, self.shape,
                               self.strides, stream=stream)
 
         self.gpu_data = gpu_data
         self.gpu_head = gpu_head
 
+    def __del__(self):
+        try:
+            self.gpu_mem.free(self.gpu_head)
+        except:
+            pass
+
 
 def from_array_like(ary, stream=0, gpu_head=None, gpu_data=None):
     "Create a DeviceNDArray object that is like ary."

numba/cuda/cudadrv/ndarray.py

@@ -3,6 +3,70 @@
 from . import devices, driver
 
 
+class ArrayHeaderManager(object):
+    """
+    Manages array header memory for reusing the allocation.
+
+    It allocates one big chunk of memory and partition it for fix sized array
+    header.  It currently stores up to 4D array header in 64-bit mode or 8D
+    array header in 32-bit mode.
+
+    This allows the small array header allocation to be reused to avoid
+    breaking asynchronous streams and avoid fragmentation of memory.
+
+    When run out of preallocated space, it automatically fallback to regular
+    allocation.
+    """
+    context_map = {}
+
+    # The number of preallocated array head
+    maxsize = 2 ** 10
+
+    # Maximum size for each array head
+    #    = 4 (ndim) * 8 (sizeof intp) * 2 (shape strides) + 8 (ptr)
+    elemsize = 72
+
+    def __new__(cls, context):
+        key = context.handle.value
+        mm = cls.context_map.get(key)
+        if mm is None:
+            mm = object.__new__(cls)
+            mm.init(context)
+            cls.context_map[key] = mm
+
+        return mm
+
+    def init(self, context):
+        self.context = context
+        self.data = self.context.memalloc(self.elemsize * self.maxsize)
+        self.queue = []
+        for i in range(self.maxsize):
+            offset = i * self.elemsize
+            mem = self.data.view(offset, offset + self.elemsize)
+            self.queue.append(mem)
+        self.allocated = set()
+
+    def allocate(self, nd):
+        arraytype = make_array_ctype(nd)
+        sizeof = ctypes.sizeof(arraytype)
+
+        # Oversized or insufficient space
+        if sizeof >= self.elemsize or not self.queue:
+            return _allocate_head(nd)
+
+        mem = self.queue.pop()
+        self.allocated.add(mem)
+        return mem
+
+    def free(self, mem):
+        if mem in self.allocated:
+            self.allocated.discard(mem)
+            self.queue.append(mem)
+
+    def __repr__(self):
+        return "<cuda managed memory %s >" % (self.context.device,)
+
+
 def make_array_ctype(ndim):
     c_intp = ctypes.c_ssize_t
 
@@ -14,8 +78,9 @@ class c_array(ctypes.Structure):
     return c_array
 
 
-def ndarray_device_allocate_head(nd):
-    "Allocate the metadata structure"
+def _allocate_head(nd):
+    """Allocate the metadata structure
+    """
     arraytype = make_array_ctype(nd)
     gpu_head = devices.get_context().memalloc(ctypes.sizeof(arraytype))
     return gpu_head
@@ -28,12 +93,6 @@ def ndarray_device_allocate_data(ary):
     return gpu_data
 
 
-def ndarray_device_transfer_data(ary, gpu_data, stream=0):
-    size = driver.host_memory_size(ary)
-    # transfer data
-    driver.host_to_device(gpu_data, ary, size, stream=stream)
-
-
 def ndarray_populate_head(gpu_head, gpu_data, shape, strides, stream=0):
     nd = len(shape)
     assert nd > 0, "0 or negative dimension"