Merge pull request #2975 from devitocodes/fuse-tasks-groups

compiler: Revamp fuse-task optoption

Author: FabioLuporini

devito/core/cpu.py

@@ -47,8 +47,8 @@ def _normalize_kwargs(cls, **kwargs):
         o['buf-async-degree'] = oo.pop('buf-async-degree', None)
         o['buf-reuse'] = oo.pop('buf-reuse', None)
 
-        # Fusion
-        o['fuse-tasks'] = oo.pop('fuse-tasks', False)
+        # Tasking
+        o['npthreads'] = oo.pop('npthreads', None)
 
         # Flops minimization
         o['cse-min-cost'] = oo.pop('cse-min-cost', cls.CSE_MIN_COST)

devito/core/gpu.py

@@ -61,8 +61,8 @@ def _normalize_kwargs(cls, **kwargs):
         o['buf-async-degree'] = oo.pop('buf-async-degree', None)
         o['buf-reuse'] = oo.pop('buf-reuse', None)
 
-        # Fusion
-        o['fuse-tasks'] = oo.pop('fuse-tasks', False)
+        # Tasking
+        o['npthreads'] = oo.pop('npthreads', None)
 
         # Flops minimization
         o['cse-min-cost'] = oo.pop('cse-min-cost', cls.CSE_MIN_COST)

devito/core/operator.py

@@ -247,6 +247,15 @@ def _check_kwargs(cls, **kwargs):
         if oo['mpi'] and oo['mpi'] not in cls.MPI_MODES:
             raise InvalidOperator(f"Unsupported MPI mode `{oo['mpi']}`")
 
+        npthreads = oo['npthreads']
+        if npthreads is not None and (
+            isinstance(npthreads, bool) or
+            not is_integer(npthreads) or npthreads <= 0
+        ):
+            raise InvalidOperator(
+                "`npthreads` must be a positive integer"
+            )
+
         if oo['cire-maxpar'] not in (False, 'basic', 'compact'):
             raise InvalidOperator("Illegal `cire-maxpar` value")
 

devito/ir/iet/nodes.py

@@ -1511,8 +1511,7 @@ class BusyWait(While):
 class SyncSpot(List):
 
     """
-    A node representing one or more synchronization operations, e.g., WaitLock,
-    withLock, etc.
+    A node coupling synchronization operations with an IET body.
     """
 
     is_SyncSpot = True

devito/ir/support/syncs.py

@@ -26,11 +26,42 @@
 
 class SyncOp(Pickable):
 
+    """
+    Metadata for a synchronization operation attached to a `Cluster` or `SyncSpot`.
+
+    Parameters
+    ----------
+    handle : object
+        The symbolic object identifying or controlling the operation, such as
+        an entry in a `Lock`. May be None when no handle is required.
+    target : AbstractFunction
+        The `Function` whose access is synchronized. For buffered data movements,
+        this is the compiler-generated buffer.
+    tindex : Expr, optional
+        The index into `target` involved in the operation.
+    function : AbstractFunction, optional
+        The original `Function` represented by a compiler-generated `target`. It
+        is the source of a `SyncCopyIn` and the destination of a `SyncCopyOut`.
+    findex : Expr, optional
+        The index into `function` corresponding to `tindex`.
+    dim : Dimension, optional
+        The `Dimension` along which `tindex` and `findex` are defined.
+    size : int, optional
+        The extent associated with the operation along `dim`. Defaults to 1.
+    origin : Indexed, optional
+        The original `Indexed` access from which the operation was derived.
+    gid : Stamp, optional
+        The `Stamp` identifying the asynchronous task group to which the operation
+        belongs.
+    """
+
     __rargs__ = ('handle', 'target')
-    __rkwargs__ = ('tindex', 'function', 'findex', 'dim', 'size', 'origin')
+    __rkwargs__ = (
+        'tindex', 'function', 'findex', 'dim', 'size', 'origin', 'gid'
+    )
 
     def __init__(self, handle, target, tindex=None, function=None, findex=None,
-                 dim=None, size=1, origin=None):
+                 dim=None, size=1, origin=None, gid=None):
         self.handle = handle
         self.target = target
 
@@ -40,6 +71,7 @@ def __init__(self, handle, target, tindex=None, function=None, findex=None,
         self.dim = dim
         self.size = size
         self.origin = origin
+        self.gid = gid
 
     def __eq__(self, other):
         return (type(self) is type(other) and
@@ -50,11 +82,13 @@ def __eq__(self, other):
                 self.findex == other.findex and
                 self.dim is other.dim and
                 self.size == other.size and
-                self.origin == other.origin)
+                self.origin == other.origin and
+                self.gid == other.gid)
 
     def __hash__(self):
         return hash((self.__class__, self.handle, self.target, self.tindex,
-                     self.function, self.findex, self.dim, self.size, self.origin))
+                     self.function, self.findex, self.dim, self.size, self.origin,
+                     self.gid))
 
     def __repr__(self):
         return f"{self.__class__.__name__}<{self.handle.name}>"

devito/passes/clusters/asynchrony.py

@@ -45,6 +45,17 @@ def memcpy_key(c):
     return task_key, memcpy_key
 
 
+def make_gid(ispace, d):
+    """
+    Make a group ID from the stamps carried by the non-trigger Intervals.
+    """
+    gids = {i.stamp for i in ispace if not i.dim._defines & d._defines}
+    if len(gids) != 1:
+        return None
+    else:
+        return gids.pop()
+
+
 @timed_pass(name='tasking')
 def tasking(clusters, key0, sregistry):
     """
@@ -147,6 +158,8 @@ def _schedule_waitlocks(self, c0, d, clusters, locks, syncs):
         return protected
 
     def _schedule_withlocks(self, c0, d, protected, locks, syncs):
+        gid = make_gid(c0.ispace, d)
+
         for target in protected:
             lock = locks[target]
 
@@ -169,7 +182,7 @@ def _schedule_withlocks(self, c0, d, protected, locks, syncs):
             for i in indices:
                 syncs[c0][d].update([
                     ReleaseLock(lock[i], target),
-                    WithLock(lock[i], target, i, function, findex, d)
+                    WithLock(lock[i], target, i, function, findex, d, gid=gid)
                 ])
 
 
@@ -274,9 +287,12 @@ def _actions_from_update_memcpy(c, d, clusters, actions, sregistry):
     guard1 = GuardBoundNext(function.indices[d], direction)
     guards = c.guards.impose(d, guard0 & guard1)
 
+    gid = make_gid(ispace, d)
+
     syncs = {d: [
         ReleaseLock(handle, target),
-        PrefetchUpdate(handle, target, tindex, function, findex, d, 1, e.rhs)
+        PrefetchUpdate(handle, target, tindex, function, findex, d,
+                       origin=e.rhs, gid=gid)
     ]}
     syncs = {**c.syncs, **syncs}
 

devito/passes/clusters/buffering.py

@@ -18,6 +18,7 @@
     timed_pass
 )
 from devito.types import Array, CustomDimension, Eq, ModuloDimension
+from devito.warnings import warn
 
 __all__ = ['buffering']
 
@@ -39,7 +40,7 @@ def buffering(clusters, key, sregistry, options, **kwargs):
         The symbol registry, to create unique names for buffers and Dimensions.
     options : dict
         The optimization options.
-        Accepted: ['buf-async-degree'].
+        Accepted: ['buf-async-degree', 'buf-reuse', 'npthreads'].
         * 'buf-async-degree': Specify the size of the buffer. By default, the
           buffer size is the minimal one, inferred from the memory accesses in
           the ``clusters`` themselves. An asynchronous degree equals to `k`
@@ -49,6 +50,9 @@ def buffering(clusters, key, sregistry, options, **kwargs):
           implemented by other passes).
         * 'buf-reuse': If True, the pass will try to reuse existing Buffers for
           different buffered Functions. By default, False.
+        * 'npthreads': Number of pthreads for asynchronous tasks. The tasks are
+          divided into this many balanced groups. By default, None, which uses
+          one pthread per task.
     **kwargs
         Additional compilation options.
         Accepted: ['opt_init_onwrite', 'opt_buffer'].
@@ -252,36 +256,41 @@ def callback(self, clusters, prefix):
                 processed.append(Cluster(expr, ispace, guards, properties, syncs))
 
         # Lift {write,read}-only buffers into separate IterationSpaces
-        if not self.options['fuse-tasks']:
-            processed = self._optimize(processed, descriptors)
+        processed = self._optimize(processed, descriptors)
 
-        if self.options['buf-reuse']:
-            init, processed = self._reuse(init, processed, descriptors)
+        # Reuse existing Buffers for buffering candidates, if requested
+        init, processed = self._reuse(init, processed, descriptors)
 
         return init + processed
 
     def _optimize(self, clusters, descriptors):
+        npthreads = self.options['npthreads']
+
+        stamps = self._make_task_groups(descriptors)
+
         for b, v in descriptors.items():
             if v.is_writeonly:
                 # `b` might be written by multiple, potentially mutually
                 # exclusive, equations. For example, two equations that have or
                 # will have complementary guards, hence only one will be
                 # executed. In such a case, we can split the equations over
                 # separate IterationSpaces
-                key0 = lambda: Stamp()
+                key0 = lambda: stamps[b] if npthreads else Stamp()  # noqa: B023
             elif v.is_readonly:
                 # `b` is read multiple times -- this could just be the case of
                 # coupled equations, so we more cautiously perform a
                 # "buffer-wise" splitting of the IterationSpaces (i.e., only
                 # relevant if there are at least two read-only buffers)
-                stamp = Stamp()
-                key0 = lambda: stamp  # noqa: B023
+                stamp_fixed = Stamp()
+                key0 = lambda: stamps[b] if npthreads else stamp_fixed  # noqa: B023
             else:
                 continue
 
             processed = []
             for c in clusters:
-                if b not in c.functions:
+                f = v.f if npthreads else b
+
+                if f not in c.functions:
                     processed.append(c)
                     continue
 
@@ -294,12 +303,57 @@ def _optimize(self, clusters, descriptors):
 
         return clusters
 
+    def _make_task_groups(self, descriptors):
+        """
+        Assign task buffers to at most `npthreads` balanced groups.
+        """
+        npthreads = self.options['npthreads']
+
+        stamps = {}
+        if not npthreads:
+            return stamps
+
+        task_sets = (
+            [b for b, v in descriptors.items() if v.is_writeonly],
+            [b for b, v in descriptors.items() if v.is_readonly],
+        )
+
+        for tasks in task_sets:
+            if not tasks:
+                continue
+
+            # Calculate the number of groups and their sizes, so that the tasks
+            # are divided into balanced groups
+            ntasks = len(tasks)
+            ngroups = min(npthreads, ntasks)
+            base, remainder = divmod(ntasks, ngroups)
+            sizes = (base + 1,) * remainder + (base,) * (ngroups - remainder)
+
+            if npthreads > ntasks:
+                warn(
+                    f"`npthreads={npthreads}` exceeds the {ntasks} available "
+                    f"tasks; using `npthreads={ntasks}` instead"
+                )
+
+            # Create a unique Stamp for each group and assign it to the tasks
+            # in that group
+            start = 0
+            for n in sizes:
+                stamp = Stamp()
+                stamps.update({b: stamp for b in tasks[start:start + n]})
+                start += n
+
+        return stamps
+
     def _reuse(self, init, clusters, descriptors):
         """
         Reuse existing Buffers for buffering candidates.
         """
         buf_reuse = self.options['buf-reuse']
 
+        if not buf_reuse:
+            return init, clusters
+
         if callable(buf_reuse):
             cbk = lambda v: [i for i in v if buf_reuse(descriptors[i])]
         else:

devito/passes/clusters/fusion.py

@@ -152,7 +152,7 @@ def __init__(self, toposort, options=None):
         options = options or {}
 
         self.toposort = toposort
-        self.fuse_tasks = options.get('fuse-tasks', False)
+        self.fuse_tasks = options.get('npthreads') is not None
 
         super().__init__()