Isolate multistage-related changes only

Author: fernanvr

devito/ir/equations/algorithms.py

@@ -6,14 +6,15 @@
 from devito.tools import (Ordering, as_tuple, flatten, filter_sorted, filter_ordered,
                           frozendict)
 from devito.types import (Dimension, Eq, IgnoreDimSort, SubDimension,
-                          ConditionalDimension)
+                          ConditionalDimension, MultiStage)
 from devito.types.array import Array
 from devito.types.basic import AbstractFunction
 from devito.types.dimension import MultiSubDimension, Thickness
 from devito.data.allocators import DataReference
 from devito.logger import warning
 
-__all__ = ['dimension_sort', 'lower_exprs', 'concretize_subdims']
+
+__all__ = ['dimension_sort', 'lower_multistage', 'lower_exprs', 'concretize_subdims']
 
 
 def dimension_sort(expr):
@@ -95,6 +96,39 @@ def handle_indexed(indexed):
     return ordering
 
 
+def lower_multistage(expressions, **kwargs):
+    """
+    Separating the multi-stage time-integrator scheme in stages:
+        * If the object is MultiStage, it creates the stages of the method.
+    """
+    return _lower_multistage(expressions, **kwargs)
+
+
+@singledispatch
+def _lower_multistage(expr, **kwargs):
+    """
+    Default handler for expressions that are not MultiStage.
+    Simply return them in a list.
+    """
+    return [expr]
+
+
+@_lower_multistage.register(MultiStage)
+def _(expr, **kwargs):
+    """
+    Specialized handler for MultiStage expressions.
+    """
+    return expr._evaluate(**kwargs)
+
+
+@_lower_multistage.register(Iterable)
+def _(exprs, **kwargs):
+    """
+    Handle iterables of expressions.
+    """
+    return sum([_lower_multistage(expr, **kwargs) for expr in exprs], [])
+
+
 def lower_exprs(expressions, subs=None, **kwargs):
     """
     Lowering an expression consists of the following passes:

devito/operations/solve.py

@@ -7,6 +7,8 @@
 from devito.finite_differences.derivative import Derivative
 from devito.tools import as_tuple
 
+from devito.types.multistage import resolve_method
+
 __all__ = ['solve', 'linsolve']
 
 
@@ -15,7 +17,7 @@ class SolveError(Exception):
     pass
 
 
-def solve(eq, target, **kwargs):
+def solve(eq, target, method = None, eq_num = 0, **kwargs):
     """
     Algebraically rearrange an Eq w.r.t. a given symbol.
 
@@ -56,9 +58,12 @@ def solve(eq, target, **kwargs):
 
     # We need to rebuild the vector/tensor as sympy.solve outputs a tuple of solutions
     if len(sols) > 1:
-        return target.new_from_mat(sols)
+        sols_temp = target.new_from_mat(sols)
     else:
-        return sols[0]
+        sols_temp = sols[0]
+
+    method = kwargs.get("method", None)
+    return sols_temp if method is None else resolve_method(method)(target, sols_temp)
 
 
 def linsolve(expr, target, **kwargs):

devito/operator/operator.py

@@ -17,7 +17,7 @@
                                InvalidOperator)
 from devito.logger import (debug, info, perf, warning, is_log_enabled_for,
                            switch_log_level)
-from devito.ir.equations import LoweredEq, lower_exprs, concretize_subdims
+from devito.ir.equations import LoweredEq, lower_multistage, lower_exprs, concretize_subdims
 from devito.ir.clusters import ClusterGroup, clusterize
 from devito.ir.iet import (Callable, CInterface, EntryFunction, DeviceFunction,
                            FindSymbols, MetaCall, derive_parameters, iet_build)
@@ -40,7 +40,6 @@
                           disk_layer)
 from devito.types.dimension import Thickness
 
-
 __all__ = ['Operator']
 
 
@@ -337,6 +336,8 @@ def _lower_exprs(cls, expressions, **kwargs):
             * Apply substitution rules;
             * Shift indices for domain alignment.
         """
+        expressions = lower_multistage(expressions, **kwargs)
+
         expand = kwargs['options'].get('expand', True)
 
         # Specialization is performed on unevaluated expressions

devito/types/__init__.py

@@ -22,3 +22,6 @@
 from .relational import *  # noqa
 from .sparse import *  # noqa
 from .tensor import *  # noqa
+
+from .multistage import *  # noqa
+from .multistage_new import *  # noqa