compiler: Pass ctx down to _map_function_on_high_bw_mem

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -563,7 +563,7 @@ class DeviceAwareDataManager(DataManager):
     def __init__(self, options=None, **kwargs):
         self.gpu_fit = options['gpu-fit']
         self.gpu_create = options['gpu-create']
-        self.pmode = options.get('place-transfers')
+        self.gpu_place_transfers = options.get('place-transfers')
 
         super().__init__(**kwargs)
 
@@ -596,7 +596,8 @@ def _map_array_on_high_bw_mem(self, site, obj, storage):
 
         storage.update(obj, site, maps=mmap, unmaps=unmap)
 
-    def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm, read_only=False):
+    def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm,
+                                     read_only=False, **kwargs):
         """
         Map a Function already defined in the host memory in to the device high
         bandwidth memory.
@@ -629,42 +630,41 @@ def _map_function_on_high_bw_mem(self, site, obj, storage, devicerm, read_only=F
         storage.update(obj, site, maps=mmap, unmaps=unmap, efuncs=efuncs)
 
     @iet_pass
-    def place_transfers(self, iet, data_movs=None, **kwargs):
+    def place_transfers(self, iet, data_movs=None, ctx=None, **kwargs):
         """
         Create a new IET with host-device data transfers. This requires mapping
         symbols to the suitable memory spaces.
         """
-        if not self.pmode:
+        if not self.gpu_place_transfers:
             return iet, {}
 
-        @singledispatch
-        def _place_transfers(iet, data_movs):
+        if not isinstance(iet, EntryFunction):
             return iet, {}
 
-        @_place_transfers.register(EntryFunction)
-        def _(iet, data_movs):
-            reads, writes = data_movs
+        reads, writes = data_movs
 
-            # Special symbol which gives user code control over data deallocations
-            devicerm = DeviceRM()
+        # Special symbol which gives user code control over data deallocations
+        devicerm = DeviceRM()
 
-            storage = Storage()
-            for i in filter_sorted(writes):
-                if i.is_Array:
-                    self._map_array_on_high_bw_mem(iet, i, storage)
-                else:
-                    self._map_function_on_high_bw_mem(iet, i, storage, devicerm)
-            for i in filter_sorted(reads - writes):
-                if i.is_Array:
-                    self._map_array_on_high_bw_mem(iet, i, storage)
-                else:
-                    self._map_function_on_high_bw_mem(iet, i, storage, devicerm, True)
-
-            iet, efuncs = self._inject_definitions(iet, storage)
+        storage = Storage()
+        for i in filter_sorted(writes):
+            if i.is_Array:
+                self._map_array_on_high_bw_mem(iet, i, storage)
+            else:
+                self._map_function_on_high_bw_mem(
+                    iet, i, storage, devicerm, ctx=ctx
+                )
+        for i in filter_sorted(reads - writes):
+            if i.is_Array:
+                self._map_array_on_high_bw_mem(iet, i, storage)
+            else:
+                self._map_function_on_high_bw_mem(
+                    iet, i, storage, devicerm, read_only=True, ctx=ctx
+                )
 
-            return iet, {'efuncs': efuncs}
+        iet, efuncs = self._inject_definitions(iet, storage)
 
-        return _place_transfers(iet, data_movs=data_movs)
+        return iet, {'efuncs': efuncs}
 
     @iet_pass
     def place_devptr(self, iet, **kwargs):