compiler: Unified Memory Allocator

devito/data/allocators.py

@@ -5,6 +5,7 @@
 import os
 import sys
 
+import cupy as cp
 import numpy as np
 import ctypes
 from ctypes.util import find_library
@@ -15,7 +16,7 @@
 
 __all__ = ['ALLOC_FLAT', 'ALLOC_NUMA_LOCAL', 'ALLOC_NUMA_ANY',
            'ALLOC_KNL_MCDRAM', 'ALLOC_KNL_DRAM', 'ALLOC_GUARD',
-           'default_allocator']
+           'CUPY_ALLOC', 'default_allocator']
 
 
 class MemoryAllocator(object):
@@ -317,6 +318,32 @@ def put_local(self):
         return self._node == 'local'
 
 
+class CupyAllocator(MemoryAllocator):
+
+    """
+    Memory allocator based on ``posix`` functions. The allocated memory is
+    aligned to page boundaries.
+    """
+
+    is_Posix = True
+
+    def __init__(self):
+        cp.cuda.set_allocator(cp.cuda.MemoryPool(cp.cuda.malloc_managed).malloc)
+
+    @classmethod
+    def initialize(cls):
+        pass
+
+
+    def _alloc_C_libcall(self, size, ctype):
+
+        mem_obj = cp.zeros(size, dtype=cp.float64)
+        return mem_obj.data.ptr, mem_obj
+
+    def free(self, c_pointer):
+        pass
+
+
 class ExternalAllocator(MemoryAllocator):
 
     """
@@ -373,6 +400,7 @@ def alloc(self, shape, dtype):
 ALLOC_KNL_MCDRAM = NumaAllocator(1)
 ALLOC_NUMA_ANY = NumaAllocator('any')
 ALLOC_NUMA_LOCAL = NumaAllocator('local')
+CUPY_ALLOC = CupyAllocator()
 
 custom_allocators = {}
 """User-defined allocators."""

devito/data/data.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from devito.data.allocators import ALLOC_FLAT
+from devito.data.allocators import ALLOC_FLAT, CUPY_ALLOC
 from devito.data.utils import *
 from devito.logger import warning
 from devito.parameters import configuration
@@ -82,7 +82,8 @@ def __del__(self):
             # Dask/Distributed context), which may (re)create a Data object
             # without going through `__array_finalize__`
             return
-        self._allocator.free(*self._memfree_args)
+        if self._allocator is not CUPY_ALLOC:
+            self._allocator.free(*self._memfree_args)
         self._memfree_args = None
 
     def __reduce__(self):

devito/passes/iet/definitions.py

@@ -20,6 +20,8 @@
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import DeviceRM, Symbol
 
+from devito.data.allocators import CUPY_ALLOC
+
 __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage']
 
 
@@ -435,6 +437,9 @@ def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm, read_only=F
         """
         mmap = self.lang._map_to(obj)
 
+        if obj._allocator is CUPY_ALLOC:
+            return
+
         if read_only is False:
             unmap = [self.lang._map_update(obj),
                      self.lang._map_release(obj, devicerm=devicerm)]