Repository: devitocodes/devito
Issue: MPIRegion pickle_args inconsistent with constructor

Hi!

We're hitting this error, since MPIRegion needs to be pickled when pickling an Operator - as needed to avoid building operators on every node of a job.

```
TypeError: __init__() missing 1 required positional argument: 'owned'
```

It appears to be caused by this:

```
class MPIRegion(CompositeObject):

    def __init__(self, name, key, arguments, owned):
        name = "%s%d" % (name, key)
        pname = "region%d" % key
...

    # Pickling support
    _pickle_args = ['name', 'arguments', 'owned']
```

I don't have an MFE to hand yet, sorry.


Changed files:
devito/symbolics/extended_sympy.py, devito/symbolics/printer.py, tests/test_pickle.py, devito/passes/iet/openmp.py, devito/mpi/routines.py, devito/core/gpu_openacc.py

diff --git a/devito/core/gpu_openacc.py b/devito/core/gpu_openacc.py
index ce3bfa668..f3baa0939 100644
--- a/devito/core/gpu_openacc.py
+++ b/devito/core/gpu_openacc.py
@@ -1,7 +1,6 @@
 from functools import partial, singledispatch
 
 import cgen as c
-from sympy import Function
 
 from devito.core.gpu_openmp import (DeviceOpenMPNoopOperator, DeviceOpenMPIteration,
                                     DeviceOmpizer, DeviceOpenMPDataManager)
@@ -12,7 +11,7 @@ from devito.logger import warning
 from devito.mpi.distributed import MPICommObject
 from devito.mpi.routines import MPICallable
 from devito.passes.iet import optimize_halospots, mpiize, hoist_prodders, iet_pass
-from devito.symbolics import Byref, Macro
+from devito.symbolics import Byref, DefFunction, Macro
 from devito.tools import as_tuple, prod, timed_pass
 from devito.types import Symbol
 
@@ -135,7 +134,7 @@ def initialize(iet, **kwargs):
             rank_init = Call('MPI_Comm_rank', [comm, Byref(rank)])
 
             ngpus = Symbol(name='ngpus')
-            call = Function('acc_get_num_devices')(device_nvidia)
+            call = DefFunction('acc_get_num_devices', device_nvidia)
             ngpus_init = LocalExpression(DummyEq(ngpus, call))
 
             devicenum = Symbol(name='devicenum')
diff --git a/devito/mpi/routines.py b/devito/mpi/routines.py
index 2f9aa0dfb..f00ef6db4 100644
--- a/devito/mpi/routines.py
+++ b/devito/mpi/routines.py
@@ -925,17 +925,19 @@ class MPIMsg(CompositeObject):
     else:
         c_mpirequest_p = type('MPI_Request', (c_void_p,), {})
 
-    def __init__(self, name, function, halos, fields=None):
+    fields = [
+        (_C_field_bufs, c_void_p),
+        (_C_field_bufg, c_void_p),
+        (_C_field_sizes, POINTER(c_int)),
+        (_C_field_rrecv, c_mpirequest_p),
+        (_C_field_rsend, c_mpirequest_p),
+    ]
+
+    def __init__(self, name, function, halos):
         self._function = function
         self._halos = halos
-        fields = (fields or []) + [
-            (MPIMsg._C_field_bufs, c_void_p),
-            (MPIMsg._C_field_bufg, c_void_p),
-            (MPIMsg._C_field_sizes, POINTER(c_int)),
-            (MPIMsg._C_field_rrecv, MPIMsg.c_mpirequest_p),
-            (MPIMsg._C_field_rsend, MPIMsg.c_mpirequest_p),
-        ]
-        super(MPIMsg, self).__init__(name, 'msg', fields)
+
+        super(MPIMsg, self).__init__(name, 'msg', self.fields)
 
         # Required for buffer allocation/deallocation before/after jumping/returning
         # to/from C-land
@@ -1009,14 +1011,12 @@ class MPIMsgEnriched(MPIMsg):
     _C_field_from = 'fromrank'
     _C_field_to = 'torank'
 
-    def __init__(self, name, function, halos):
-        fields = [
-            (MPIMsgEnriched._C_field_ofss, POINTER(c_int)),
-            (MPIMsgEnriched._C_field_ofsg, POINTER(c_int)),
-            (MPIMsgEnriched._C_field_from, c_int),
-            (MPIMsgEnriched._C_field_to, c_int)
-        ]
-        super(MPIMsgEnriched, self).__init__(name, function, halos, fields)
+    fields = MPIMsg.fields + [
+        (_C_field_ofss, POINTER(c_int)),
+        (_C_field_ofsg, POINTER(c_int)),
+        (_C_field_from, c_int),
+        (_C_field_to, c_int)
+    ]
 
     def _arg_defaults(self, alias=None):
         super(MPIMsgEnriched, self)._arg_defaults(alias)
@@ -1054,15 +1054,17 @@ class MPIMsgEnriched(MPIMsg):
 
 class MPIRegion(CompositeObject):
 
-    def __init__(self, name, key, arguments, owned):
-        name = "%s%d" % (name, key)
-        pname = "region%d" % key
-
+    def __init__(self, prefix, key, arguments, owned):
+        self._prefix = prefix
+        self._key = key
         self._owned = owned
 
         # Sorting for deterministic codegen
         self._arguments = sorted(arguments, key=lambda i: i.name)
 
+        name = "%s%d" % (prefix, key)
+        pname = "region%d" % key
+
         fields = []
         for i in self.arguments:
             if i.is_Dimension:
@@ -1083,6 +1085,14 @@ class MPIRegion(CompositeObject):
     def arguments(self):
         return self._arguments
 
+    @property
+    def prefix(self):
+        return self._prefix
+
+    @property
+    def key(self):
+        return self._key
+
     @property
     def owned(self):
         return self._owned
@@ -1108,4 +1118,4 @@ class MPIRegion(CompositeObject):
         return values
 
     # Pickling support
-    _pickle_args = ['name', 'arguments', 'owned']
+    _pickle_args = ['prefix', 'key', 'arguments', 'owned']
diff --git a/devito/passes/iet/openmp.py b/devito/passes/iet/openmp.py
index 11e31c8ce..4ff6f75b6 100644
--- a/devito/passes/iet/openmp.py
+++ b/devito/passes/iet/openmp.py
@@ -3,13 +3,13 @@ import os
 
 import numpy as np
 import cgen as c
-from sympy import Function, Or, Max, Not
+from sympy import Or, Max, Not
 
 from devito.ir import (DummyEq, Conditional, Block, Expression, ExpressionBundle, List,
                        Prodder, Iteration, While, FindSymbols, FindNodes, Return,
                        COLLAPSED, VECTORIZED, Transformer, IsPerfectIteration,
                        retrieve_iteration_tree, filter_iterations)
-from devito.symbolics import CondEq, INT
+from devito.symbolics import CondEq, DefFunction, INT
 from devito.parameters import configuration
 from devito.passes.iet.engine import iet_pass
 from devito.tools import as_tuple, is_integer, prod
@@ -193,12 +193,12 @@ class ThreadedProdder(Conditional, Prodder):
 
     def __init__(self, prodder):
         # Atomic-ize any single-thread Prodders in the parallel tree
-        condition = CondEq(Function('omp_get_thread_num')(), 0)
+        condition = CondEq(DefFunction('omp_get_thread_num'), 0)
 
         # Prod within a while loop until all communications have completed
         # In other words, the thread delegated to prodding is entrapped for as long
         # as it's required
-        prod_until = Not(Function(prodder.name)(*[i.name for i in prodder.arguments]))
+        prod_until = Not(DefFunction(prodder.name, [i.name for i in prodder.arguments]))
         then_body = List(header=c.Comment('Entrap thread until comms have completed'),
                          body=While(prod_until))
 
diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py
index 6be0b88be..35a94b162 100644
--- a/devito/symbolics/extended_sympy.py
+++ b/devito/symbolics/extended_sympy.py
@@ -11,7 +11,8 @@ from devito.tools import Pickable, as_tuple, is_integer
 
 __all__ = ['CondEq', 'CondNe', 'IntDiv', 'FunctionFromPointer', 'FieldFromPointer',
            'FieldFromComposite', 'ListInitializer', 'Byref', 'IndexedPointer',
-           'Macro', 'Literal', 'INT', 'FLOAT', 'DOUBLE', 'FLOOR', 'cast_mapper']
+           'DefFunction', 'Macro', 'Literal', 'INT', 'FLOAT', 'DOUBLE', 'FLOOR',
+           'cast_mapper']
 
 
 class CondEq(sympy.Eq):
@@ -256,7 +257,7 @@ class Byref(sympy.Expr, Pickable):
     __reduce_ex__ = Pickable.__reduce_ex__
 
 
-class IndexedPointer(sympy.Expr):
+class IndexedPointer(sympy.Expr, Pickable):
 
     """
     Symbolic representation of the C notation ``symbol[...]``
@@ -293,6 +294,40 @@ class IndexedPointer(sympy.Expr):
     __reduce_ex__ = Pickable.__reduce_ex__
 
 
+class DefFunction(Function, Pickable):
+
+    """
+    A definitely-defined sympy.Function, to work around:
+
+        https://github.com/sympy/sympy/issues/4297
+    """
+
+    def __new__(cls, name, arguments=None):
+        arguments = as_tuple(arguments)
+        obj = Function.__new__(cls, name, *arguments)
+        obj._name = name
+        obj._arguments = arguments
+        return obj
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def arguments(self):
+        return self._arguments
+
+    def __str__(self):
+        return "%s(%s)" % (self.name, ', '.join(str(i) for i in self.arguments))
+
+    __repr__ = __str__
+
+    # Pickling support
+    _pickle_args = ['name']
+    _pickle_kwargs = ['arguments']
+    __reduce_ex__ = Pickable.__reduce_ex__
+
+
 class Macro(sympy.Symbol):
 
     """
diff --git a/devito/symbolics/printer.py b/devito/symbolics/printer.py
index 8fb24304b..deae77a53 100644
--- a/devito/symbolics/printer.py
+++ b/devito/symbolics/printer.py
@@ -117,6 +117,7 @@ class CodePrinter(C99CodePrinter):
 
     _print_Byref = _print_IntDiv
     _print_IndexedPointer = _print_IntDiv
+    _print_DefFunction = _print_IntDiv
 
     def _print_TrigonometricFunction(self, expr):
         func_name = str(expr.func)
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index ab1658858..d3e1fec4f 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -1,16 +1,19 @@
 import pytest
 import numpy as np
-from sympy import Symbol
+from sympy import Symbol, Min
 import pickle
 
 from conftest import skipif
 from devito import (Constant, Eq, Function, TimeFunction, SparseFunction, Grid,
                     Dimension, SubDimension, ConditionalDimension, IncrDimension,
                     TimeDimension, SteppingDimension, Operator, ShiftedDimension)
-from devito.mpi.routines import MPIStatusObject, MPIRequestObject
+from devito.data import LEFT, OWNED
+from devito.mpi.halo_scheme import Halo
+from devito.mpi.routines import (MPIStatusObject, MPIMsgEnriched, MPIRequestObject,
+                                 MPIRegion)
 from devito.operator.profiling import Timer
 from devito.types import Symbol as dSymbol, Scalar
-from devito.symbolics import IntDiv, ListInitializer, FunctionFromPointer
+from devito.symbolics import IntDiv, ListInitializer, FunctionFromPointer, DefFunction
 from examples.seismic import (demo_model, AcquisitionGeometry,
                               TimeAxis, RickerSource, Receiver)
 
@@ -221,6 +224,12 @@ def test_symbolics():
     new_li = pickle.loads(pkl_li)
     assert li == new_li
 
+    df = DefFunction('f', ['a', 1, 2])
+    pkl_df = pickle.dumps(df)
+    new_df = pickle.loads(pkl_df)
+    assert df == new_df
+    assert df.arguments == new_df.arguments
+
 
 def test_timers():
     """Pickling for Timers used in Operators for C-level profiling."""
@@ -325,8 +334,9 @@ def test_operator_timefunction_w_preallocation():
 @skipif(['nompi'])
 @pytest.mark.parallel(mode=[1])
 def test_mpi_objects():
-    # Neighbours
     grid = Grid(shape=(4, 4, 4))
+
+    # Neighbours
     obj = grid.distributor._obj_neighborhood
     pkl_obj = pickle.dumps(obj)
     new_obj = pickle.loads(pkl_obj)
@@ -357,11 +367,51 @@ def test_mpi_objects():
 
 
 @skipif(['nompi'])
-@pytest.mark.parallel(mode=[1])
+@pytest.mark.parallel(mode=[(1, 'full')])
+def test_mpi_fullmode_objects():
+    grid = Grid(shape=(4, 4, 4))
+    x, y, _ = grid.dimensions
+
+    # Message
+    f = Function(name='f', grid=grid)
+    obj = MPIMsgEnriched('msg', f, [Halo(x, LEFT)])
+    pkl_obj = pickle.dumps(obj)
+    new_obj = pickle.loads(pkl_obj)
+    assert obj.name == new_obj.name
+    assert obj.function.name == new_obj.function.name
+    assert all(obj.function.dimensions[i].name == new_obj.function.dimensions[i].name
+               for i in range(grid.dim))
+    assert new_obj.function.dimensions[0] is new_obj.halos[0].dim
+
+    # Region
+    x_m, x_M = x.symbolic_min, x.symbolic_max
+    y_m, y_M = y.symbolic_min, y.symbolic_max
+    obj = MPIRegion('reg', 1, [y, x],
+                    [(((x, OWNED, LEFT),), {x: (x_m, Min(x_M, x_m))}),
+                     (((y, OWNED, LEFT),), {y: (y_m, Min(y_M, y_m))})])
+    pkl_obj = pickle.dumps(obj)
+    new_obj = pickle.loads(pkl_obj)
+    assert obj.prefix == new_obj.prefix
+    assert obj.key == new_obj.key
+    assert obj.name == new_obj.name
+    assert len(new_obj.arguments) == 2
+    assert all(d0.name == d1.name for d0, d1 in zip(obj.arguments, new_obj.arguments))
+    assert all(new_obj.arguments[i] is new_obj.owned[i][0][0][0]  # `x` and `y`
+               for i in range(2))
+    assert new_obj.owned[0][0][0][1] is new_obj.owned[1][0][0][1]  # `OWNED`
+    assert new_obj.owned[0][0][0][2] is new_obj.owned[1][0][0][2]  # `LEFT`
+    for n, i in enumerate(new_obj.owned):
+        d, v = list(i[1].items())[0]
+        assert d is new_obj.arguments[n]
+        assert v[0] is d.symbolic_min
+        assert v[1] == Min(d.symbolic_max, d.symbolic_min)
+
+
+@skipif(['nompi'])
+@pytest.mark.parallel(mode=[(1, 'basic'), (1, 'full')])
 def test_mpi_operator():
     grid = Grid(shape=(4,))
     f = TimeFunction(name='f', grid=grid)
-    g = TimeFunction(name='g', grid=grid)
 
     # Using `sum` creates a stencil in `x`, which in turn will
     # trigger the generation of code for MPI halo exchange
@@ -373,6 +423,8 @@ def test_mpi_operator():
 
     assert str(op) == str(new_op)
 
+    new_grid = new_op.input[0].grid
+    g = TimeFunction(name='g', grid=new_grid)
     new_op.apply(time=2, f=g)
     assert np.all(f.data[0] == [2., 3., 3., 3.])
     assert np.all(f.data[1] == [3., 6., 7., 7.])