compiler: Move and enhance FunctionMap

Author: FabioLuporini

devito/types/misc.py

@@ -381,6 +381,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

@@ -20,13 +20,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__ = [
     'Barrier',
     'DeviceID',
     'DeviceRM',
-    'FunctionMap',
     'Lock',
     'NPThreads',
     'NThreads',
@@ -386,19 +384,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

@@ -22,8 +22,9 @@
 )
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
 from devito.types import (
-    Array, CustomDimension, FunctionMap, LocalObject, Pointer, Symbol
+    Array, CustomDimension, LocalObject, Pointer, Symbol
 )
+from devito.types.misc import FunctionMap
 
 
 @pytest.fixture
@@ -322,7 +323,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
@@ -342,10 +343,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():