compiler: Postpone weights unevaluation

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, unevaluate  # noqa, sigh
-        weights = tuple(unevaluate(pow_to_mul(sympy.sympify(i))) for i in weights)
+        from devito.symbolics import pow_to_mul  # noqa, sigh
+        weights = tuple(pow_to_mul(sympy.sympify(i)) for i in weights)
 
         kwargs['scope'] = kwargs.get('scope', 'stack')
         kwargs['initvalue'] = weights

devito/passes/iet/definitions.py

@@ -16,7 +16,7 @@
 from devito.passes.iet.engine import iet_pass
 from devito.passes.iet.langbase import LangBB
 from devito.symbolics import (Byref, DefFunction, FieldFromPointer, IndexedPointer,
-                              SizeOf, VOID, pow_to_mul)
+                              SizeOf, VOID, pow_to_mul, unevaluate)
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap,
                           DeviceRM, Eq, Symbol)
@@ -119,7 +119,7 @@ def _alloc_array_on_global_mem(self, site, obj, storage):
 
         # Create input array
         name = '%s_init' % obj.name
-        initvalue = np.array([pow_to_mul(i) for i in obj.initvalue])
+        initvalue = np.array([unevaluate(pow_to_mul(i)) for i in obj.initvalue])
         src = Array(name=name, dtype=obj.dtype, dimensions=obj.dimensions,
                     space='host', scope='stack', initvalue=initvalue)
 

tests/test_pickle.py

@@ -1,7 +1,7 @@
 import ctypes
 import pickle as pickle0
-import cloudpickle as pickle1
 
+import cloudpickle as pickle1
 import pytest
 import numpy as np
 from sympy import Symbol
@@ -12,15 +12,14 @@
                     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, ComponentAccess, StencilDimension)
+                          FIndexed, ComponentAccess)
 from devito.types.basic import BoundSymbol, AbstractSymbol
 from devito.tools import EnrichedTuple
 from devito.symbolics import (IntDiv, ListInitializer, FieldFromPointer,
@@ -431,25 +430,6 @@ def test_component_access(self, pickle):
         assert new_ca.index == 1
         assert new_ca.function.name == f.name
 
-    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')