compiler: Add optoption for collecting derivatives

Author: FabioLuporini

devito/core/cpu.py

@@ -76,6 +76,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Code generation options for derivatives
         o['expand'] = oo.pop('expand', cls.EXPAND)
+        o['deriv-collect'] = oo.pop('deriv-collect', cls.DERIV_COLLECT)
         o['deriv-schedule'] = oo.pop('deriv-schedule', cls.DERIV_SCHEDULE)
         o['deriv-unroll'] = oo.pop('deriv-unroll', False)
 
@@ -151,7 +152,7 @@ class Cpu64AdvOperator(Cpu64OperatorMixin, CoreOperator):
     @classmethod
     @timed_pass(name='specializing.DSL')
     def _specialize_dsl(cls, expressions, **kwargs):
-        expressions = collect_derivatives(expressions)
+        expressions = collect_derivatives(expressions, **kwargs)
 
         return expressions
 
@@ -254,7 +255,7 @@ class Cpu64CustomOperator(Cpu64OperatorMixin, CustomOperator):
     @classmethod
     def _make_dsl_passes_mapper(cls, **kwargs):
         return {
-            'collect-derivs': collect_derivatives,
+            'deriv-collect': collect_derivatives,
         }
 
     @classmethod
@@ -309,7 +310,7 @@ def _make_iet_passes_mapper(cls, **kwargs):
 
     _known_passes = (
         # DSL
-        'collect-derivs',
+        'deriv-collect',
         # Expressions
         'buffering',
         # Clusters

devito/core/gpu.py

@@ -89,6 +89,7 @@ def _normalize_kwargs(cls, **kwargs):
 
         # Code generation options for derivatives
         o['expand'] = oo.pop('expand', cls.EXPAND)
+        o['deriv-collect'] = oo.pop('deriv-collect', cls.DERIV_COLLECT)
         o['deriv-schedule'] = oo.pop('deriv-schedule', cls.DERIV_SCHEDULE)
         o['deriv-unroll'] = oo.pop('deriv-unroll', False)
 
@@ -189,7 +190,7 @@ class DeviceAdvOperator(DeviceOperatorMixin, CoreOperator):
     @classmethod
     @timed_pass(name='specializing.DSL')
     def _specialize_dsl(cls, expressions, **kwargs):
-        expressions = collect_derivatives(expressions)
+        expressions = collect_derivatives(expressions, **kwargs)
 
         return expressions
 
@@ -281,7 +282,7 @@ class DeviceCustomOperator(DeviceOperatorMixin, CustomOperator):
     @classmethod
     def _make_dsl_passes_mapper(cls, **kwargs):
         return {
-            'collect-derivs': collect_derivatives,
+            'deriv-collect': collect_derivatives,
         }
 
     @classmethod
@@ -331,7 +332,7 @@ def _make_iet_passes_mapper(cls, **kwargs):
 
     _known_passes = (
         # DSL
-        'collect-derivs',
+        'deriv-collect',
         # Expressions
         'buffering',
         # Clusters

devito/core/operator.py

@@ -123,6 +123,12 @@ class BasicOperator(Operator):
     finite-difference derivatives.
     """
 
+    DERIV_COLLECT = True
+    """
+    Factorize finite-difference derivatives exploiting the linearity of the FD
+    operators.
+    """
+
     DERIV_SCHEDULE = 'basic'
     """
     The schedule to use for the computation of finite-difference derivatives.
@@ -296,7 +302,7 @@ def _specialize_dsl(cls, expressions, **kwargs):
         # Call passes
         for i in passes:
             try:
-                expressions = passes_mapper[i](expressions)
+                expressions = passes_mapper[i](expressions, **kwargs)
             except KeyError:
                 pass
 

devito/passes/equations/linearity.py

@@ -11,12 +11,23 @@
 
 
 @timed_pass()
-def collect_derivatives(expressions):
+def collect_derivatives(expressions, options=None, **kwargs):
     """
     Exploit linearity of finite-differences to collect `Derivative`'s of
     same type. This may help CIRE creating fewer temporaries while catching
     larger redundant sub-expressions.
     """
+    deriv_collect = options['deriv-collect']
+    if not deriv_collect:
+        return expressions
+
+    return _collect_derivatives(expressions)
+
+
+def _collect_derivatives(expressions):
+    """
+    Carry out the bulk of the work for `collect_derivatives`.
+    """
     processed = []
     for e in expressions:
         # Track type and number of nested Derivatives

tests/test_dse.py

@@ -2099,7 +2099,7 @@ def test_sum_of_nested_derivatives(self, expr, exp_arrays, exp_ops):
         eqn = Eq(v.forward, eval(expr))
 
         op0 = Operator(eqn, opt=('noop', {'openmp': True}))
-        op1 = Operator(eqn, opt=('collect-derivs', 'cire-sops', {'openmp': True}))
+        op1 = Operator(eqn, opt=('deriv-collect', 'cire-sops', {'openmp': True}))
         op2 = Operator(eqn, opt=('cire-sops', {'openmp': True}))
         op3 = Operator(eqn, opt=('advanced', {'openmp': True}))
 

tests/test_lower_exprs.py

@@ -6,7 +6,9 @@
 from devito.finite_differences import Derivative
 from devito.finite_differences.differentiable import diff2sympy
 from devito.ir.equations import LoweredEq
-from devito.passes.equations.linearity import collect_derivatives
+from devito.passes.equations.linearity import (
+    _collect_derivatives as collect_derivatives
+)
 from devito.tools import timed_region
 
 
@@ -34,8 +36,9 @@ def test_nocollection_if_diff_dims(self):
 
             # Since all Function are time dependent, there should be no collection
             # and produce the same result as with the pre evaluated expression
-            expr = Operator._lower_exprs([eq], options={})[0]
-            expr2 = Operator._lower_exprs([eq.evaluate], options={})[0]
+            options = {'deriv-collect': True}
+            expr = Operator._lower_exprs([eq], options=options)[0]
+            expr2 = Operator._lower_exprs([eq.evaluate], options=options)[0]
 
         assert expr == expr2
 
@@ -45,7 +48,7 @@ def test_numeric_constant(self):
         u = TimeFunction(name="u", grid=grid, space_order=4, time_order=2)
 
         eq = Eq(u.forward, u.dx.dx + 0.3*u.dy.dx)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert len(leq.find(Derivative)) == 3
 
@@ -56,7 +59,7 @@ def test_symbolic_constant(self):
         u = TimeFunction(name="u", grid=grid, space_order=4, time_order=2)
 
         eq = Eq(u.forward, u.dx.dx + dt**0.2*u.dy.dx)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert len(leq.find(Derivative)) == 3
 
@@ -69,7 +72,7 @@ def test_symbolic_constant_times_add(self):
 
         eq = Eq(u.forward, u.laplace + dt**0.2*u.biharmonic(1/f))
 
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert len(eq.rhs.args) == 3
         assert len(leq.rhs.args) == 2
@@ -86,7 +89,7 @@ def test_solve(self):
 
         pde = u.dt2 - (u.dx.dx + u.dy.dy) - u.dx.dy
         eq = Eq(u.forward, solve(pde, u.forward))
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert len(eq.rhs.find(Derivative)) == 5
         assert len(leq.rhs.find(Derivative)) == 4
@@ -99,7 +102,7 @@ def test_nocollection_if_unworthy(self):
         u = TimeFunction(name="u", grid=grid)
 
         eq = Eq(u.forward, (0.4 + dt)*(u.dx + u.dy))
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq == leq
 
@@ -112,7 +115,7 @@ def test_pull_and_collect(self):
         v = TimeFunction(name="v", grid=grid)
 
         eq = Eq(u.forward, ((0.4 + dt)*u.dx + 0.3)*hx + v.dx)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq != leq
         args = leq.rhs.args
@@ -129,7 +132,7 @@ def test_pull_and_collect_nested(self):
         v = TimeFunction(name="v", grid=grid, space_order=2)
 
         eq = Eq(u.forward, (((0.4 + dt)*u.dx + 0.3)*hx + v.dx).dy + (0.2 + hy)*v.dy)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq != leq
         assert leq.rhs == ((v + hx*(0.4 + dt)*u).dx + 0.3*hx + (0.2 + hy)*v).dy
@@ -143,7 +146,7 @@ def test_pull_and_collect_nested_v2(self):
         v = TimeFunction(name="v", grid=grid, space_order=2)
 
         eq = Eq(u.forward, ((0.4 + dt*(hy + 1. + hx*hy))*u.dx + 0.3)*hx + v.dx)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq != leq
         assert leq.rhs == 0.3*hx + (hx*(0.4 + dt*(hy + 1. + hx*hy))*u + v).dx
@@ -158,7 +161,7 @@ def test_pull_and_collect_nested_v3(self):
         v = TimeFunction(name="v", grid=grid, space_order=2)
 
         eq = Eq(u.forward, 0.4 + a*(hx + dt*(u.dx + v.dx)))
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq != leq
         assert leq.rhs == 0.4 + a*(hx + (dt*u + dt*v).dx)
@@ -172,7 +175,7 @@ def test_nocollection_subdims(self):
         f = Function(name='f', grid=grid)
 
         eq = Eq(u.forward, u.dx + 0.2*f[xi, yi]*v.dx)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq == leq
 
@@ -184,7 +187,7 @@ def test_nocollection_staggered(self):
         v = TimeFunction(name="v", grid=grid, staggered=x)
 
         eq = Eq(u.forward, u.dx + v.dx)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq == leq
 
@@ -195,13 +198,13 @@ def test_nocollection_mixed_order(self):
 
         # First case is obvious...
         eq = Eq(u.forward, u.dx2 + u.dx.dy + 1.)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq == leq
 
         # y-derivative should not get collected!
         eq = Eq(u.forward, u.dy2 + u.dx.dy + 1.)
-        leq = collect_derivatives.func([eq])[0]
+        leq = collect_derivatives([eq])[0]
 
         assert eq == leq