compiler: Add ThreadArrive and TensorMove

Author: FabioLuporini

devito/ir/clusters/cluster.py

@@ -16,7 +16,8 @@
 from devito.symbolics import estimate_cost
 from devito.tools import as_tuple, filter_ordered, flatten, infer_dtype
 from devito.types import (
-    CriticalRegion, Fence, ThreadCommit, ThreadPoolSync, ThreadWait, WeakFence
+    CriticalRegion, Fence, Indexed, TensorMove, ThreadArrive, ThreadCommit,
+    ThreadPoolSync, ThreadWait, WeakFence
 )
 
 __all__ = ["Cluster", "ClusterGroup"]
@@ -310,14 +311,43 @@ def is_critical_region(self):
     def is_thread_pool_sync(self):
         return self._is_type(ThreadPoolSync)
 
+    @cached_property
+    def is_shm_write(self):
+        return all(w._mem_shared for w in self.scope.writes)
+
     @cached_property
     def is_thread_commit(self):
         return self._is_type(ThreadCommit)
 
+    @cached_property
+    def is_thread_arrive(self):
+        return self._is_type(ThreadArrive)
+
     @cached_property
     def is_thread_wait(self):
         return self._is_type(ThreadWait)
 
+    @cached_property
+    def is_word_move(self):
+        return (self._is_type(Indexed) and
+                all(e.rhs.function._mem_heap for e in self.exprs))
+
+    @cached_property
+    def is_tensor_move(self):
+        return self._is_type(TensorMove)
+
+    @cached_property
+    def is_word_move_to_mem_shared(self):
+        return self.is_word_move and self.is_shm_write
+
+    @cached_property
+    def is_tensor_move_to_mem_shared(self):
+        return self.is_tensor_move and self.is_shm_write
+
+    @cached_property
+    def is_glb_load_to_mem_shared(self):
+        return self.is_word_move_to_mem_shared or self.is_tensor_move_to_mem_shared
+
     @cached_property
     def is_async(self):
         """
@@ -557,6 +587,10 @@ def dspace(self):
     def is_halo_touch(self):
         return all(i.is_halo_touch for i in self)
 
+    @cached_property
+    def is_glb_load_to_mem_shared(self):
+        return all(i.is_glb_load_to_mem_shared for i in self)
+
     @cached_property
     def dtype(self):
         """

devito/ir/support/properties.py

@@ -97,11 +97,6 @@ def __init__(self, name, val=None):
 A Dimension along which prefetching is feasible and beneficial.
 """
 
-PREFETCHABLE_SHM = Property('prefetchable-shm')
-"""
-A Dimension along which shared-memory prefetching is feasible and beneficial.
-"""
-
 INIT_CORE_SHM = Property('init-core-shm')
 """
 A Dimension along which the shared-memory CORE data region is initialized.
@@ -190,32 +185,6 @@ def update_properties(properties, exprs):
     if not exprs:
         return properties
 
-    # Auto-detect prefetchable Dimensions
-    dims = set()
-    flag = False
-    for e in as_tuple(exprs):
-        w, r = e.args
-
-        # Ensure it's in the form `Indexed = Indexed`
-        try:
-            wf, rf = w.function, r.function
-        except AttributeError:
-            break
-
-        if not rf or not wf._mem_shared:
-            break
-        dims.update({d.parent for d in wf.dimensions if d.parent in properties})
-
-        if not rf._mem_heap:
-            break
-    else:
-        flag = True
-
-    if flag:
-        properties = properties.prefetchable_shm(dims)
-    else:
-        properties = properties.drop(properties=PREFETCHABLE_SHM)
-
     # Remove properties that are trivially incompatible with `exprs`
     if not all(e.lhs.function._mem_shared for e in as_tuple(exprs)):
         drop = {INIT_CORE_SHM, INIT_HALO_LEFT_SHM, INIT_HALO_RIGHT_SHM}
@@ -284,9 +253,6 @@ def prefetchable(self, dims, v=PREFETCHABLE):
             m[d] = self.get(d, set()) | {v}
         return Properties(m)
 
-    def prefetchable_shm(self, dims):
-        return self.prefetchable(dims, PREFETCHABLE_SHM)
-
     def block(self, dims, kind='default'):
         if kind == 'default':
             p = TILABLE
@@ -357,9 +323,6 @@ def _is_property_any(self, dims, v):
     def is_prefetchable(self, dims=None, v=PREFETCHABLE):
         return self._is_property_any(dims, PREFETCHABLE)
 
-    def is_prefetchable_shm(self, dims=None):
-        return self._is_property_any(dims, PREFETCHABLE_SHM)
-
     def is_core_init(self, dims=None):
         return self._is_property_any(dims, INIT_CORE_SHM)
 

devito/passes/clusters/misc.py

@@ -232,7 +232,7 @@ def _key(self, c):
         weak.append(c.properties.is_core_init())
 
         # Prefetchable Clusters should get merged, if possible
-        weak.append(c.properties.is_prefetchable_shm())
+        weak.append(c.is_glb_load_to_mem_shared)
 
         # Promoting adjacency of IndexDerivatives will maximize their reuse
         weak.append(any(search(c.exprs, IndexDerivative)))

devito/types/parallel.py

@@ -11,10 +11,11 @@
 from functools import cached_property
 
 import numpy as np
+from sympy import Expr
 
 from devito.exceptions import InvalidArgument
 from devito.parameters import configuration
-from devito.symbolics import search
+from devito.symbolics import Reserved, Terminal, search
 from devito.tools import as_list, as_tuple, is_integer
 from devito.types.array import Array, ArrayObject
 from devito.types.basic import Scalar, Symbol
@@ -35,7 +36,9 @@
     'QueueID',
     'SharedData',
     'TBArray',
+    'TensorMove',
     'ThreadArray',
+    'ThreadArrive',
     'ThreadCommit',
     'ThreadID',
     'ThreadPoolSync',
@@ -365,12 +368,24 @@ class ThreadCommit(Fence):
     pass
 
 
+class ThreadArrive(Fence):
+
+    """
+    A generic arrive operation for a single thread, typically used to signal
+    the arrival at a certain point through a suitable synchronization object.
+    """
+
+    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.
+    with other threads over:
+
+        * a memory operation issued by a prior ThreadCommit operation.
+        * the consumption of a shared resource via a ThreadArrive operation.
     """
 
     pass
@@ -386,3 +401,18 @@ def __init_finalize__(self, *args, **kwargs):
         kwargs['liveness'] = 'eager'
 
         super().__init_finalize__(*args, **kwargs)
+
+
+class TensorMove(Expr, Reserved, Terminal):
+
+    """
+    Represent the LOAD/STORE of a multi-dimensional block of data from/to a higher
+    level of the memory hierarchy
+    """
+
+    func = Reserved._rebuild
+
+    def _ccode(self, printer):
+        return str(self)
+
+    _sympystr = _ccode