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
 

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
 

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.

devito/passes/equations/linearity.py

@@ -11,12 +11,16 @@
 
 
 @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
+
     processed = []
     for e in expressions:
         # Track type and number of nested Derivatives