Merge pull request #2794 from devitocodes/reuse-shmas-for-wdr

compiler: Add special types for device operations

Author: mloubout

devito/ir/clusters/cluster.py

@@ -15,7 +15,9 @@
 from devito.mpi.reduction_scheme import DistReduce
 from devito.symbolics import estimate_cost
 from devito.tools import as_tuple, filter_ordered, flatten, infer_dtype
-from devito.types import Fence, WeakFence, CriticalRegion
+from devito.types import (
+    Fence, WeakFence, CriticalRegion, ThreadPoolSync, ThreadCommit, ThreadWait
+)
 
 __all__ = ["Cluster", "ClusterGroup"]
 
@@ -262,26 +264,40 @@ def is_wild(self):
                 self.is_weak_fence or
                 self.is_critical_region)
 
+    def _is_type(self, cls):
+        return self.exprs and all(isinstance(e.rhs, cls) for e in self.exprs)
+
     @cached_property
     def is_halo_touch(self):
-        return self.exprs and all(isinstance(e.rhs, HaloTouch) for e in self.exprs)
+        return self._is_type(HaloTouch)
 
     @cached_property
     def is_dist_reduce(self):
-        return self.exprs and all(isinstance(e.rhs, DistReduce) for e in self.exprs)
+        return self._is_type(DistReduce)
 
     @cached_property
     def is_fence(self):
-        return (self.exprs and all(isinstance(e.rhs, Fence) for e in self.exprs) or
-                self.is_critical_region)
+        return self._is_type(Fence) or self.is_critical_region
 
     @cached_property
     def is_weak_fence(self):
-        return self.exprs and all(isinstance(e.rhs, WeakFence) for e in self.exprs)
+        return self._is_type(WeakFence)
 
     @cached_property
     def is_critical_region(self):
-        return self.exprs and all(isinstance(e.rhs, CriticalRegion) for e in self.exprs)
+        return self._is_type(CriticalRegion)
+
+    @cached_property
+    def is_thread_pool_sync(self):
+        return self._is_type(ThreadPoolSync)
+
+    @cached_property
+    def is_thread_commit(self):
+        return self._is_type(ThreadCommit)
+
+    @cached_property
+    def is_thread_wait(self):
+        return self._is_type(ThreadWait)
 
     @cached_property
     def is_async(self):

devito/types/parallel.py

@@ -22,7 +22,8 @@
 
 __all__ = ['NThreads', 'NThreadsNested', 'NThreadsNonaffine', 'NThreadsBase',
            'DeviceID', 'ThreadID', 'Lock', 'ThreadArray', 'PThreadArray',
-           'SharedData', 'NPThreads', 'DeviceRM', 'QueueID', 'Barrier', 'TBArray']
+           'SharedData', 'NPThreads', 'DeviceRM', 'QueueID', 'Barrier', 'TBArray',
+           'ThreadPoolSync', 'ThreadCommit', 'ThreadWait']
 
 
 class NThreadsAbstract(Scalar):
@@ -321,6 +322,38 @@ class Barrier(Fence):
     pass
 
 
+class ThreadPoolSync(Barrier):
+
+    """
+    A generic synchronization barrier for a pool of threads.
+    """
+
+    pass
+
+
+class ThreadCommit(Fence):
+
+    """
+    A generic commit operation for a single thread, typically used to issue
+    a memory operation at a specific program point, which requires the special
+    treatment that all Fence subclasses provide (i.e., to avoid being reshuffled
+    around by optimization passes).
+    """
+
+    pass
+
+
+class ThreadWait(Fence):
+
+    """
+    A generic wait operation for a single thread, typically used to synchronize
+    after a memory operation issued at a specific program point with a
+    ThreadCommit operation.
+    """
+
+    pass
+
+
 class TBArray(Array):
 
     """