compiler: Add FunctionMap type

Author: FabioLuporini

devito/types/parallel.py

@@ -20,11 +20,13 @@
 from devito.types.basic import Scalar, Symbol
 from devito.types.dimension import CustomDimension
 from devito.types.misc import Fence, VolatileInt
+from devito.types.object import LocalObject
 
 __all__ = [
     'Barrier',
     'DeviceID',
     'DeviceRM',
+    'FunctionMap',
     'Lock',
     'NPThreads',
     'NThreads',
@@ -384,3 +386,19 @@ def __init_finalize__(self, *args, **kwargs):
         kwargs['liveness'] = 'eager'
 
         super().__init_finalize__(*args, **kwargs)
+
+
+class FunctionMap(LocalObject):
+
+    """
+    Wrap a Function in a LocalObject.
+    """
+
+    __rargs__ = ('name', 'tensor')
+
+    def __init__(self, name, tensor, **kwargs):
+        super().__init__(name, **kwargs)
+        self.tensor = tensor
+
+    def _hashable_content(self):
+        return super()._hashable_content() + (self.tensor,)

tests/test_iet.py

@@ -17,10 +17,13 @@
 from devito.passes.iet.engine import Graph
 from devito.passes.iet.languages.C import CDataManager
 from devito.symbolics import (
-    FLOAT, Byref, Class, FieldFromComposite, InlineIf, Macro, String
+    FLOAT, Byref, Class, FieldFromComposite, InlineIf, ListInitializer,
+    Macro, SizeOf, String
 )
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
-from devito.types import Array, CustomDimension, LocalObject, Pointer, Symbol
+from devito.types import (
+    Array, CustomDimension, FunctionMap, LocalObject, Pointer, Symbol
+)
 
 
 @pytest.fixture
@@ -299,6 +302,52 @@ def _C_free(self):
 }"""
 
 
+def test_make_cuda_tensor_map():
+
+    class CUTensorMap(FunctionMap):
+
+        dtype = CustomDtype('CUtensorMap')
+
+        @property
+        def _C_init(self):
+            symsizes = list(reversed(self.tensor.symbolic_shape))
+            sizeof_dtype = SizeOf(self.tensor.dmap._C_typedata)
+
+            sizes = ListInitializer(symsizes)
+            strides = ListInitializer([
+                np.prod(symsizes[:i])*sizeof_dtype for i in range(1, len(symsizes))
+            ])
+
+            arguments = [
+                Byref(self),
+                Macro('CU_TENSOR_MAP_DATA_TYPE_FLOAT32'),
+                4, self.tensor.dmap, sizes, strides,
+                ]
+            call = Call('cuTensorMapEncodeTiled', arguments)
+
+            return call
+
+    grid = Grid(shape=(10, 10, 10))
+
+    u = TimeFunction(name='u', grid=grid)
+
+    tmap = CUTensorMap('tmap', u)
+
+    iet = Call('foo', tmap)
+    iet = ElementalFunction('foo', iet, parameters=())
+    dm = CDataManager(sregistry=None)
+    iet = CDataManager.place_definitions.__wrapped__(dm, iet)[0]
+
+    assert str(iet) == """\
+static void foo()
+{
+  CUtensorMap tmap;
+  cuTensorMapEncodeTiled(&(tmap),CU_TENSOR_MAP_DATA_TYPE_FLOAT32,4,d_u,{u_vec->size[3], u_vec->size[2], u_vec->size[1], u_vec->size[0]},{sizeof(float)*u_vec->size[3], sizeof(float)*u_vec->size[2]*u_vec->size[3], sizeof(float)*u_vec->size[1]*u_vec->size[2]*u_vec->size[3]});
+
+  foo(tmap);
+}"""
+
+
 def test_cpp_local_object():
     """
     Test C++ support for LocalObjects.