compiler: Add cire-minmem optoption

Author: FabioLuporini

devito/core/cpu.py

@@ -61,6 +61,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['cire-maxpar'] = oo.pop('cire-maxpar', False)
         o['cire-ftemps'] = oo.pop('cire-ftemps', False)
         o['cire-mingain'] = oo.pop('cire-mingain', cls.CIRE_MINGAIN)
+        o['cire-minmem'] = oo.pop('cire-minmem', cls.CIRE_MINMEM)
         o['cire-schedule'] = oo.pop('cire-schedule', cls.CIRE_SCHEDULE)
 
         # Shared-memory parallelism

devito/core/gpu.py

@@ -68,6 +68,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['cire-maxpar'] = oo.pop('cire-maxpar', True)
         o['cire-ftemps'] = oo.pop('cire-ftemps', False)
         o['cire-mingain'] = oo.pop('cire-mingain', cls.CIRE_MINGAIN)
+        o['cire-minmem'] = oo.pop('cire-minmem', cls.CIRE_MINMEM)
         o['cire-schedule'] = oo.pop('cire-schedule', cls.CIRE_SCHEDULE)
 
         # GPU parallelism

devito/core/operator.py

@@ -69,6 +69,14 @@ class BasicOperator(Operator):
     intensity of the generated kernel.
     """
 
+    CIRE_MINMEM = True
+    """
+    Minimize memory consumption when allocating temporaries for CIRE-optimized
+    expressions. This may come at the cost of slighly worse performance due to
+    the potential need for extra registers to hold a greater number of support
+    variables (e.g., strides).
+    """
+
     SCALAR_MIN_TYPE = np.float16
     """
     Minimum datatype for a scalar arising from a common sub-expression or CIRE temp.

devito/passes/clusters/aliases.py

@@ -14,7 +14,7 @@
                        vmax, vmin)
 from devito.passes.clusters.cse import _cse
 from devito.symbolics import (Uxmapper, estimate_cost, search, reuse_if_untouched,
-                              uxreplace, sympy_dtype)
+                              retrieve_functions, uxreplace, sympy_dtype)
 from devito.tools import (Stamp, as_mapper, as_tuple, flatten, frozendict,
                           is_integer, generator, split, timed_pass)
 from devito.types import (Eq, Symbol, Temp, TempArray, TempFunction,
@@ -113,6 +113,7 @@ def __init__(self, sregistry, options, platform):
         self.opt_rotate = options['cire-rotate']
         self.opt_ftemps = options['cire-ftemps']
         self.opt_mingain = options['cire-mingain']
+        self.opt_minmem = options['cire-minmem']
         self.opt_min_dtype = options['scalar-min-type']
         self.opt_multisubdomain = True
 
@@ -143,7 +144,8 @@ def _aliases_from_clusters(self, clusters, exclude, meta):
 
         # Schedule -> [Clusters]_k
         processed, subs = lower_schedule(schedule, meta, self.sregistry,
-                                         self.opt_ftemps, self.opt_min_dtype)
+                                         self.opt_ftemps, self.opt_min_dtype,
+                                         self.opt_minmem)
 
         # [Clusters]_k -> [Clusters]_k (optimization)
         if self.opt_multisubdomain:
@@ -831,11 +833,12 @@ def optimize_schedule_rotations(schedule, sregistry):
     return schedule.rebuild(*processed, rmapper=rmapper)
 
 
-def lower_schedule(schedule, meta, sregistry, ftemps, min_dtype):
+def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
+                   opt_minmem):
     """
     Turn a Schedule into a sequence of Clusters.
     """
-    if ftemps:
+    if opt_ftemps:
         make = TempFunction
     else:
         # Typical case -- the user does *not* "see" the CIRE-created temporaries
@@ -865,8 +868,26 @@ def lower_schedule(schedule, meta, sregistry, ftemps, min_dtype):
             dimensions = [d.parent if d.is_AbstractSub else d
                           for d in writeto.itdims]
 
-            # The halo must be set according to the size of `writeto`
-            halo = [(abs(i.lower), abs(i.upper)) for i in writeto]
+            # The minimum halo required along each Dimension depends on `writeto`.
+            # The user might suggest to go more relaxed about this via `opt_minmem`,
+            # in which case we extend the halo based on the surrounding
+            # Functions to minimize support variables such as strides etc
+            halo = {i.dim: Size(abs(i.lower), abs(i.upper)) for i in writeto}
+
+            if opt_minmem:
+                functions = []
+            else:
+                functions = retrieve_functions(pivot)
+
+            for f in functions:
+                for d, h0 in list(halo.items()):
+                    try:
+                        h1 = f._size_halo[d]
+                    except KeyError:
+                        continue
+                    halo[d] = Size(max(h0.left, h1.left), max(h0.right, h1.right))
+
+            halo = tuple(halo.values())
 
             # The indices used to write into the Array
             indices = []
@@ -889,7 +910,7 @@ def lower_schedule(schedule, meta, sregistry, ftemps, min_dtype):
             # Degenerate case: scalar expression
             assert writeto.size == 0
 
-            dtype = sympy_dtype(pivot, base=meta.dtype, smin=min_dtype)
+            dtype = sympy_dtype(pivot, base=meta.dtype, smin=opt_min_dtype)
             obj = Temp(name=name, dtype=dtype)
             expression = Eq(obj, uxreplace(pivot, subs))
 
@@ -1037,6 +1058,9 @@ def pick_best(variants):
 # Utilities
 
 
+Size = namedtuple('Size', 'left right')
+
+
 class Group(tuple):
 
     """

devito/passes/iet/linearization.py

@@ -67,7 +67,7 @@ def key1(f, d):
 
         * False if not statically linearizable, that is not linearizable via
           constant symbolic sizes and strides;
-        * A 3-tuple `(Dimension, halo size, grid)` otherwise.
+        * A 3-tuple `(Dimension, halo size, pad dtype)` otherwise.
     """
     if f.is_regular:
         # For paddable objects the following holds:

tests/test_linearize.py

@@ -659,3 +659,29 @@ def test_int64_array(order):
     else:
         long = '(long)'
         assert f'({2*order} + {long}y_size)*({2*order} + {long}x_size))' in str(op)
+
+
+def test_cire_n_strides():
+    grid = Grid(shape=(4, 4, 4))
+
+    u = TimeFunction(name='u', grid=grid, space_order=8)
+    u1 = TimeFunction(name='u', grid=grid, space_order=8)
+
+    eqn = Eq(u.forward, u.dy.dx + u.dy.dy + u.dy.dz + 1.)
+
+    op0 = Operator(eqn, opt=('advanced', {'linearize': False, 'cire-mingain': 0}))
+    op1 = Operator(eqn, opt=('advanced', {'linearize': True, 'cire-mingain': 0}))
+    op2 = Operator(eqn, opt=('advanced', {'linearize': True,
+                                          'cire-mingain': 0,
+                                          'cire-minmem': False}))
+
+    # Check generated code
+    assert 'uL0' in str(op1)
+    assert len(op1.body.strides) == 11
+    assert 'uL0' in str(op2)
+    assert len(op2.body.strides) == 9  # Fewer size/stride vars thx to cire-minmem
+
+    op0.apply(time_M=10)
+    op2.apply(time_M=10, u=u1)
+
+    assert np.all(u.data == u1.data)