compiler: Move and enhance FunctionMap

Author: FabioLuporini

devito/types/misc.py

@@ -345,6 +345,30 @@ def closing(self):
 """
 
 
+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,)
+
+    @property
+    def free_symbols(self):
+        """
+        The free symbols of a FunctionMap are the free symbols of the
+        underlying Function.
+        """
+        return super().free_symbols | {self.tensor}
+
+
 # *** C/CXX support types
 
 size_t = CustomDtype('size_t')

devito/types/parallel.py

@@ -19,12 +19,10 @@
 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__ = ['NThreads', 'NThreadsNested', 'NThreadsNonaffine', 'NThreadsBase',
            'DeviceID', 'ThreadID', 'Lock', 'ThreadArray', 'PThreadArray',
-           'SharedData', 'NPThreads', 'DeviceRM', 'QueueID', 'Barrier', 'TBArray',
-           'FunctionMap']
+           'SharedData', 'NPThreads', 'DeviceRM', 'QueueID', 'Barrier', 'TBArray']
 
 
 class NThreadsAbstract(Scalar):
@@ -333,19 +331,3 @@ 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,7 +17,8 @@
 from devito.symbolics import (Byref, FieldFromComposite, InlineIf, Macro, Class,
                               FLOAT, ListInitializer, SizeOf)
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
-from devito.types import Array, LocalObject, Symbol, FunctionMap
+from devito.types import Array, LocalObject, Symbol
+from devito.types.misc import FunctionMap
 
 
 @pytest.fixture
@@ -316,7 +317,7 @@ def _C_init(self):
                 Byref(self),
                 Macro('CU_TENSOR_MAP_DATA_TYPE_FLOAT32'),
                 4, self.tensor.dmap, sizes, strides,
-                ]
+            ]
             call = Call('cuTensorMapEncodeTiled', arguments)
 
             return call
@@ -336,10 +337,10 @@ def _C_init(self):
 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]});
+  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);
-}"""
+}"""  # noqa
 
 
 def test_cpp_local_object():