Merge pull request #2674 from devitocodes/min-buffer-mem

compiler: Add buf-reuse opt-option and rcompile caching

Author: mloubout

devito/core/cpu.py

@@ -34,6 +34,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Buffering
         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)

devito/core/gpu.py

@@ -44,6 +44,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Buffering
         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)

devito/operator/operator.py

@@ -25,9 +25,10 @@
 from devito.operator.registry import operator_selector
 from devito.mpi import MPI
 from devito.parameters import configuration
-from devito.passes import (Graph, lower_index_derivatives, generate_implicit,
-                           generate_macros, minimize_symbols, unevaluate,
-                           error_mapper, is_on_device, lower_dtypes)
+from devito.passes import (
+    Graph, lower_index_derivatives, generate_implicit, generate_macros,
+    minimize_symbols, unevaluate, error_mapper, is_on_device, lower_dtypes
+)
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple,
                           flatten, filter_sorted, frozendict, is_integer,
@@ -1150,6 +1151,35 @@ def __setstate__(self, state):
         )
 
 
+# *** Recursive compilation ("rcompile") machinery
+
+
+class RCompiles(CacheInstances):
+
+    """
+    A cache for abstract Callables obtained from lowering expressions.
+    Here, "abstract Callable" means that any user-level symbolic object appearing
+    in the input expressions is replaced by a corresponding abstract object.
+    """
+
+    _instance_cache_size = None
+
+    def __init__(self, exprs, cls):
+        self.exprs = exprs
+        self.cls = cls
+
+        # NOTE: Constructed lazily at `__call__` time because `**kwargs` is
+        # unhashable for historical reasons (e.g., Compiler objects are mutable,
+        # though in practice they are unique per Operator, so only "locally"
+        # mutable)
+        self._output = None
+
+    def compile(self, **kwargs):
+        if self._output is None:
+            self._output = self.cls._lower(self.exprs, **kwargs)
+        return self._output
+
+
 # Default action (perform or bypass) for selected compilation passes upon
 # recursive compilation
 # NOTE: it may not only be pointless to apply the following passes recursively
@@ -1167,6 +1197,7 @@ def rcompile(expressions, kwargs, options, target=None):
     """
     Perform recursive compilation on an ordered sequence of symbolic expressions.
     """
+    expressions = as_tuple(expressions)
     options = {**options, **rcompile_registry}
 
     if target is None:
@@ -1181,10 +1212,14 @@ def rcompile(expressions, kwargs, options, target=None):
     # Recursive profiling not supported -- would be a complete mess
     kwargs.pop('profiler', None)
 
-    return cls._lower(expressions, **kwargs)
+    # Recursive compilation is expensive, so we cache the result because sometimes
+    # it is called multiple times for the same input
+    compiled = RCompiles(expressions, cls).compile(**kwargs)
+
+    return compiled
 
 
-# Misc helpers
+# *** Misc helpers
 
 
 IRs = namedtuple('IRs', 'expressions clusters stree uiet iet')

devito/passes/clusters/buffering.py

@@ -44,12 +44,17 @@ def buffering(clusters, key, sregistry, options, **kwargs):
           ModuloDimensions. This might help relieving the synchronization
           overhead when asynchronous operations are used (these are however
           implemented by other passes).
+        * 'buf-reuse': If True, the pass will try to reuse existing Buffers for
+          different buffered Functions. By default, False.
     **kwargs
         Additional compilation options.
         Accepted: ['opt_init_onwrite', 'opt_buffer'].
         * 'opt_init_onwrite': By default, a written buffer does not trigger the
         generation of an initializing Cluster. With `opt_init_onwrite=True`,
         instead, the buffer gets initialized to zero.
+        * 'opt_reuse': A callback that takes a buffering candidate `bf` as input
+        and returns True if the pass can reuse pre-existing Buffers for
+        buffering `bf`, which would otherwise default to False.
         * 'opt_buffer': A callback that takes a buffering candidate as input
         and returns a buffer, which would otherwise default to an Array.
 
@@ -98,6 +103,7 @@ def key(f):
     options.update({
         'buf-init-onwrite': init_onwrite,
         'buf-callback': kwargs.get('opt_buffer'),
+        'buf-reuse': kwargs.get('opt_reuse', options['buf-reuse']),
     })
 
     # Escape hatch to selectively disable buffering
@@ -246,10 +252,13 @@ def callback(self, clusters, prefix):
                 processed.append(Cluster(expr, ispace, guards, properties, syncs))
 
         # Lift {write,read}-only buffers into separate IterationSpaces
-        if self.options['fuse-tasks']:
-            return init + processed
-        else:
-            return init + self._optimize(processed, descriptors)
+        if not self.options['fuse-tasks']:
+            processed = self._optimize(processed, descriptors)
+
+        if self.options['buf-reuse']:
+            init, processed = self._reuse(init, processed, descriptors)
+
+        return init + processed
 
     def _optimize(self, clusters, descriptors):
         for b, v in descriptors.items():
@@ -285,6 +294,48 @@ def _optimize(self, clusters, descriptors):
 
         return clusters
 
+    def _reuse(self, init, clusters, descriptors):
+        """
+        Reuse existing Buffers for buffering candidates.
+        """
+        buf_reuse = self.options['buf-reuse']
+
+        if callable(buf_reuse):
+            cbk = lambda v: [i for i in v if buf_reuse(descriptors[i].f)]
+        else:
+            cbk = lambda v: v
+
+        mapper = as_mapper(descriptors, key=lambda b: b._signature)
+        mapper = {k: cbk(v) for k, v in mapper.items() if cbk(v)}
+
+        subs = {}
+        drop = set()
+        for reusable in mapper.values():
+            retain = reusable.pop(0)
+            drop.update(reusable)
+
+            name = self.sregistry.make_name(prefix='r')
+            b = retain.func(name=name)
+
+            for i in (retain, *reusable):
+                subs.update({i: b, i.indexed: b.indexed})
+
+        processed = []
+        for c in init:
+            if set(c.scope.writes) & drop:
+                continue
+
+            exprs = [uxreplace(e, subs) for e in c.exprs]
+            processed.append(c.rebuild(exprs=exprs))
+        init = processed
+
+        processed = []
+        for c in clusters:
+            exprs = [uxreplace(e, subs) for e in c.exprs]
+            processed.append(c.rebuild(exprs=exprs))
+
+        return init, processed
+
 
 Map = namedtuple('Map', 'b f')
 

tests/test_buffering.py

@@ -724,3 +724,28 @@ def test_stencil_issue_1915_v2(subdomain):
     op1.apply(time_M=nt-2, u=u1)
 
     assert np.all(u.data == u1.data)
+
+
+def test_buffer_reuse():
+    nt = 10
+    grid = Grid(shape=(4, 4))
+
+    u = TimeFunction(name='u', grid=grid)
+    usave = TimeFunction(name='usave', grid=grid, save=nt)
+    vsave = TimeFunction(name='vsave', grid=grid, save=nt)
+
+    eqns = [Eq(u.forward, u + 1),
+            Eq(usave, u.forward),
+            Eq(vsave, u.forward + 1)]
+
+    op = Operator(eqns, opt=('buffering', {'buf-reuse': True}))
+
+    # Check generated code
+    assert len(retrieve_iteration_tree(op)) == 5
+    buffers = [i for i in FindSymbols().visit(op) if i.is_Array and i._mem_heap]
+    assert len(buffers) == 1
+
+    op.apply(time_M=nt-1)
+
+    assert all(np.all(usave.data[i-1] == i) for i in range(1, nt + 1))
+    assert all(np.all(vsave.data[i-1] == i + 1) for i in range(1, nt + 1))