compiler: Fix Weights reconstruction

Author: FabioLuporini

devito/finite_differences/differentiable.py

@@ -746,8 +746,8 @@ def __init_finalize__(self, *args, **kwargs):
         assert isinstance(weights, (list, tuple, np.ndarray))
 
         # Normalize `weights`
-        from devito.symbolics import pow_to_mul  # noqa, sigh
-        weights = tuple(pow_to_mul(sympy.sympify(i)) for i in weights)
+        from devito.symbolics import pow_to_mul, unevaluate  # noqa, sigh
+        weights = tuple(unevaluate(pow_to_mul(sympy.sympify(i))) for i in weights)
 
         kwargs['scope'] = kwargs.get('scope', 'stack')
         kwargs['initvalue'] = weights

devito/passes/clusters/unevaluate.py

@@ -1,8 +1,5 @@
-import sympy
-
 from devito.ir import cluster_pass
-from devito.symbolics import reuse_if_untouched, q_leaf
-from devito.symbolics.unevaluation import Add, Mul, Pow
+from devito.symbolics import unevaluate as _unevaluate
 
 __all__ = ['unevaluate']
 
@@ -12,22 +9,3 @@ def unevaluate(cluster):
     exprs = [_unevaluate(e) for e in cluster.exprs]
 
     return cluster.rebuild(exprs=exprs)
-
-
-mapper = {
-    sympy.Add: Add,
-    sympy.Mul: Mul,
-    sympy.Pow: Pow
-}
-
-
-def _unevaluate(expr):
-    if q_leaf(expr):
-        return expr
-
-    args = [_unevaluate(a) for a in expr.args]
-
-    try:
-        return mapper[expr.func](*args)
-    except KeyError:
-        return reuse_if_untouched(expr, args)

devito/symbolics/manipulation.py

@@ -13,7 +13,9 @@
 from devito.symbolics.extended_sympy import DefFunction, rfunc
 from devito.symbolics.queries import q_leaf
 from devito.symbolics.search import retrieve_indexed, retrieve_functions
-from devito.symbolics.unevaluation import Mul as UMul
+from devito.symbolics.unevaluation import (
+    Add as UnevalAdd, Mul as UnevalMul, Pow as UnevalPow
+)
 from devito.tools import as_list, as_tuple, flatten, split, transitive_closure
 from devito.types.basic import Basic, Indexed
 from devito.types.array import ComponentAccess
@@ -22,7 +24,7 @@
 
 __all__ = ['xreplace_indices', 'pow_to_mul', 'indexify', 'subs_op_args',
            'normalize_args', 'uxreplace', 'Uxmapper', 'subs_if_composite',
-           'reuse_if_untouched', 'evalrel', 'flatten_args']
+           'reuse_if_untouched', 'evalrel', 'flatten_args', 'unevaluate']
 
 
 def uxreplace(expr, rule):
@@ -338,7 +340,7 @@ def pow_to_mul(expr):
             # but at least we traverse the base looking for other Pows
             return expr.func(pow_to_mul(base), exp, evaluate=False)
         elif exp > 0:
-            return UMul(*[pow_to_mul(base)]*int(exp), evaluate=False)
+            return UnevalMul(*[pow_to_mul(base)]*int(exp), evaluate=False)
         elif exp < 0:
             # Reciprocal powers become inverse of the negative power
             # for example Pow(expr, -2) becomes Pow(expr * expr, -1)
@@ -502,3 +504,18 @@ def evalrel(func=min, input=None, assumptions=None):
     except TypeError:
         pass
     return rfunc(func, *input)
+
+
+uneval_mapper = {Add: UnevalAdd, Mul: UnevalMul, Pow: UnevalPow}
+
+
+def unevaluate(expr):
+    if q_leaf(expr):
+        return expr
+
+    args = [unevaluate(a) for a in expr.args]
+
+    try:
+        return uneval_mapper[expr.func](*args)
+    except KeyError:
+        return reuse_if_untouched(expr, args)

tests/test_pickle.py

@@ -12,14 +12,15 @@
                     PrecomputedSparseTimeFunction, SubDomain)
 from devito.ir import Backward, Forward, GuardFactor, GuardBound, GuardBoundNext
 from devito.data import LEFT, OWNED
+from devito.finite_differences.differentiable import Weights
 from devito.finite_differences.tools import direct, transpose, left, right, centered
 from devito.mpi.halo_scheme import Halo
 from devito.mpi.routines import (MPIStatusObject, MPIMsgEnriched, MPIRequestObject,
                                  MPIRegion)
 from devito.types import (Array, CustomDimension, Symbol as dSymbol, Scalar,
                           PointerArray, Lock, PThreadArray, SharedData, Timer,
                           DeviceID, NPThreads, ThreadID, TempFunction, Indirection,
-                          FIndexed)
+                          FIndexed, StencilDimension)
 from devito.types.basic import BoundSymbol, AbstractSymbol
 from devito.tools import EnrichedTuple
 from devito.symbolics import (IntDiv, ListInitializer, FieldFromPointer,
@@ -416,6 +417,25 @@ def test_findexed(self, pickle):
         assert new_fi.indices == (x+1, y, z-2)
         assert new_fi.strides_map == fi.strides_map
 
+    def test_weights_to_array(self, pickle):
+        grid = Grid(shape=(3, 3, 3))
+        x, y, z = grid.dimensions
+        h_x = x.spacing
+
+        i = StencilDimension('i0', 0, 2)
+        w = Weights(name='w0', dimensions=i,
+                    initvalue=[1/(h_x**2), 2/(h_x**2), 3/(h_x**2)])
+        a = Array(name='w0', dimensions=w.dimensions, initvalue=w.initvalue,
+                  scope='stack')
+
+        pkl_a = pickle.dumps(a)
+        new_a = pickle.loads(pkl_a)
+
+        # Weights optimizes `initvalue` by turning pows into muls. This test checks
+        # that the optimization is correctly carried over to the pickled object
+        # (in practice, the optimized expressions must have been frozen)
+        assert a.initvalue == new_a.initvalue
+
     def test_symbolics(self, pickle):
         a = Symbol('a')