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']
+           'ALLOC_CUPY', 'default_allocator']
 
 
 class MemoryAllocator(object):
@@ -317,6 +318,28 @@ def put_local(self):
         return self._node == 'local'
 
 
+class CupyAllocator(MemoryAllocator):
+
+    """
+    Memory allocator based on Unified Memory concept. The allocation is made using Cupy.
+    """
+
+    def __init__(self):
+        self.mempool = cp.cuda.MemoryPool(cp.cuda.malloc_managed)
+        cp.cuda.set_allocator(self.mempool.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, _):
+        self.mempool.free_all_blocks()
+
+
 class ExternalAllocator(MemoryAllocator):
 
     """
@@ -373,6 +396,7 @@ def alloc(self, shape, dtype):
 ALLOC_KNL_MCDRAM = NumaAllocator(1)
 ALLOC_NUMA_ANY = NumaAllocator('any')
 ALLOC_NUMA_LOCAL = NumaAllocator('local')
+ALLOC_CUPY = CupyAllocator()
 
 custom_allocators = {}
 """User-defined allocators."""

tests/test_data.py

@@ -1,8 +1,9 @@
 import pytest
 import numpy as np
+import cupy as cp
 
 from devito import (Grid, Function, TimeFunction, SparseTimeFunction, Dimension, # noqa
-                    Eq, Operator, ALLOC_GUARD, ALLOC_FLAT, configuration, switchconfig)
+                    Eq, Operator, ALLOC_GUARD, ALLOC_FLAT, ALLOC_CUPY, configuration, switchconfig)
 from devito.data import LEFT, RIGHT, Decomposition, loc_data_idx, convert_index
 from devito.tools import as_tuple
 from devito.types import Scalar
@@ -206,6 +207,26 @@ def test_indexing_into_sparse(self):
         sf.data[1:-1, 0] = np.arange(8)
         assert np.all(sf.data[1:-1, 0] == np.arange(8))
 
+    def test_uma_allocation(self):
+        """
+        Test Unified Memory allocation.
+        """
+        nt = 5
+        grid = Grid(shape=(4, 4, 4))
+
+        u = Function(name='u', grid=grid, allocator=ALLOC_CUPY )
+        u.data[:] = 5
+        address = u.data.ctypes.data
+        pointerAttr = cp.cuda.runtime.pointerGetAttributes(address)
+        assert pointerAttr.devicePointer == pointerAttr.hostPointer
+
+        v = TimeFunction(name='v', grid=grid, save=nt, allocator=ALLOC_CUPY )
+        v.data[:] = 5
+        address = v.data.ctypes.data
+        pointerAttr = cp.cuda.runtime.pointerGetAttributes(address)
+        assert pointerAttr.devicePointer == pointerAttr.hostPointer
+
+
 
 class TestLocDataIDX(object):
     """