compiler: Move and enhance FunctionMap

Author: FabioLuporini

devito/types/misc.py

@@ -358,6 +358,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,11 @@
 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',
-           'ThreadPoolSync', 'ThreadCommit', 'ThreadWait', 'FunctionMap']
+           'ThreadPoolSync', 'ThreadCommit', 'ThreadWait']
 
 
 class NThreadsAbstract(Scalar):
@@ -365,19 +364,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

@@ -20,8 +20,9 @@
                               String, ListInitializer, SizeOf, FLOAT)
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
 from devito.types import (
-    CustomDimension, Array, LocalObject, Symbol, Pointer, FunctionMap
+    CustomDimension, Array, LocalObject, Symbol, Pointer
 )
+from devito.types.misc import FunctionMap
 
 
 @pytest.fixture
@@ -320,7 +321,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
@@ -340,10 +341,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():