ckp: Checkpointing update

devito/__init__.py

@@ -22,6 +22,7 @@
 from devito.builtins import *  # noqa
 from devito.data.allocators import *  # noqa
 from devito.mpi import MPI  # noqa
+from devito.checkpointing import DevitoCheckpoint, CheckpointOperator  # noqa
 
 # Imports required to initialize Devito
 from devito.arch import compiler_registry, platform_registry

devito/checkpointing/checkpoint.py

@@ -0,0 +1,97 @@
+from pyrevolve import Checkpoint, Operator
+from devito import TimeFunction
+from devito.tools import flatten
+
+
+class CheckpointOperator(Operator):
+    """Devito's concrete implementation of the ABC pyrevolve.Operator. This class wraps
+       devito.Operator so it conforms to the pyRevolve API. pyRevolve will call apply
+       with arguments t_start and t_end. Devito calls these arguments time_m and time_M
+       so the following dict is used to perform the translations between different names.
+
+       Parameters
+       ----------
+       op : Operator
+            devito.Operator object that this object will wrap.
+       args : dict
+            If devito.Operator.apply() expects any arguments, they can be provided
+            here to be cached. Any calls to CheckpointOperator.apply() will
+            automatically include these cached arguments in the call to the
+            underlying devito.Operator.apply().
+    """
+    t_arg_names = {'t_start': 'time_m', 't_end': 'time_M'}
+
+    def __init__(self, op, **kwargs):
+        self.op = op
+        self.args = self.op._prepare_arguments(**kwargs)
+        self.start_offset = self.args[self.t_arg_names['t_start']]
+
+    def _prepare_args(self, t_start, t_end):
+        args = self.args.copy()
+        args[self.t_arg_names['t_start']] = t_start + self.start_offset
+        args[self.t_arg_names['t_end']] = t_end - 1 + self.start_offset
+        return args
+
+    def apply(self, t_start, t_end):
+        """ If the devito operator requires some extra arguments in the call to apply
+            they can be stored in the args property of this object so pyRevolve calls
+            pyRevolve.Operator.apply() without caring about these extra arguments while
+            this method passes them on correctly to devito.Operator
+        """
+        # Build the arguments list to invoke the kernel function
+        args = self._prepare_args(t_start, t_end)
+        # Invoke kernel function with args
+        arg_values = [args[p.name] for p in self.op.parameters]
+        self.op.cfunction(*arg_values)
+
+
+class DevitoCheckpoint(Checkpoint):
+    """Devito's concrete implementation of the Checkpoint abstract base class provided by
+       pyRevolve. Holds a list of symbol objects that hold data.
+    """
+    def __init__(self, objects):
+        """Intialise a checkpoint object. Upon initialisation, a checkpoint
+        stores only a reference to the objects that are passed into it."""
+        assert(all(isinstance(o, TimeFunction) for o in objects))
+        dtypes = set([o.dtype for o in objects])
+        assert(len(dtypes) == 1)
+        self._dtype = dtypes.pop()
+        self.objects = objects
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    def get_data(self, timestep):
+        data = flatten([get_symbol_data(s, timestep) for s in self.objects])
+        return data
+
+    def get_data_location(self, timestep):
+        return self.get_data(timestep)
+
+    @property
+    def size(self):
+        """The memory consumption of the data contained in a checkpoint."""
+        return sum([int((o.size_allocated/(o.time_order+1))*o.time_order)
+                   for o in self.objects])
+
+    def save(*args):
+        raise RuntimeError("Invalid method called. Did you check your version" +
+                           " of pyrevolve?")
+
+    def load(*args):
+        raise RuntimeError("Invalid method called. Did you check your version" +
+                           " of pyrevolve?")
+
+
+def get_symbol_data(symbol, timestep):
+    timestep += symbol.time_order - 1
+    ptrs = []
+    for i in range(symbol.time_order):
+        # Use `._data`, instead of `.data`, as `.data` is a view of the DOMAIN
+        # data region which is non-contiguous in memory. The performance hit from
+        # dealing with non-contiguous memory is so big (introduces >1 copy), it's
+        # better to checkpoint unneccesarry stuff to get a contiguous chunk of memory.
+        ptr = symbol._data[timestep - i, :, :]
+        ptrs.append(ptr)
+    return ptrs

examples/checkpointing/checkpoint.py

@@ -1,98 +1,7 @@
-from pyrevolve import Checkpoint, Operator
-from devito import TimeFunction
-from devito.tools import flatten
+from devito import warning
 
+warning("""The location of Devito's checkpointing has changed. This location will be
+           deprecated soon. Please change your imports to 'from devito import
+           DevitoCheckpoint, CheckpointOperato'""")
 
-class CheckpointOperator(Operator):
-    """Devito's concrete implementation of the ABC pyrevolve.Operator. This class wraps
-       devito.Operator so it conforms to the pyRevolve API. pyRevolve will call apply
-       with arguments t_start and t_end. Devito calls these arguments t_s and t_e so
-       the following dict is used to perform the translations between different names.
-
-       Parameters
-       ----------
-       op : Operator
-            devito.Operator object that this object will wrap.
-       args : dict
-            If devito.Operator.apply() expects any arguments, they can be provided
-            here to be cached. Any calls to CheckpointOperator.apply() will
-            automatically include these cached arguments in the call to the
-            underlying devito.Operator.apply().
-    """
-    t_arg_names = {'t_start': 'time_m', 't_end': 'time_M'}
-
-    def __init__(self, op, **kwargs):
-        self.op = op
-        self.args = kwargs
-        op_default_args = self.op._prepare_arguments(**kwargs)
-        self.start_offset = op_default_args[self.t_arg_names['t_start']]
-
-    def _prepare_args(self, t_start, t_end):
-        args = self.args.copy()
-        args[self.t_arg_names['t_start']] = t_start + self.start_offset
-        args[self.t_arg_names['t_end']] = t_end - 1 + self.start_offset
-        return args
-
-    def apply(self, t_start, t_end):
-        """ If the devito operator requires some extra arguments in the call to apply
-            they can be stored in the args property of this object so pyRevolve calls
-            pyRevolve.Operator.apply() without caring about these extra arguments while
-            this method passes them on correctly to devito.Operator
-        """
-        # Build the arguments list to invoke the kernel function
-        args = self.op.arguments(**self._prepare_args(t_start, t_end))
-        # Invoke kernel function with args
-        arg_values = [args[p.name] for p in self.op.parameters]
-        self.op.cfunction(*arg_values)
-
-
-class DevitoCheckpoint(Checkpoint):
-    """Devito's concrete implementation of the Checkpoint abstract base class provided by
-       pyRevolve. Holds a list of symbol objects that hold data.
-    """
-    def __init__(self, objects):
-        """Intialise a checkpoint object. Upon initialisation, a checkpoint
-        stores only a reference to the objects that are passed into it."""
-        assert(all(isinstance(o, TimeFunction) for o in objects))
-        dtypes = set([o.dtype for o in objects])
-        assert(len(dtypes) == 1)
-        self._dtype = dtypes.pop()
-        self.objects = objects
-
-    @property
-    def dtype(self):
-        return self._dtype
-
-    def get_data(self, timestep):
-        data = flatten([get_symbol_data(s, timestep) for s in self.objects])
-        return data
-
-    def get_data_location(self, timestep):
-        return self.get_data(timestep)
-
-    @property
-    def size(self):
-        """The memory consumption of the data contained in a checkpoint."""
-        return sum([int((o.size_allocated/(o.time_order+1))*o.time_order)
-                   for o in self.objects])
-
-    def save(*args):
-        raise RuntimeError("Invalid method called. Did you check your version" +
-                           " of pyrevolve?")
-
-    def load(*args):
-        raise RuntimeError("Invalid method called. Did you check your version" +
-                           " of pyrevolve?")
-
-
-def get_symbol_data(symbol, timestep):
-    timestep += symbol.time_order - 1
-    ptrs = []
-    for i in range(symbol.time_order):
-        # Use `._data`, instead of `.data`, as `.data` is a view of the DOMAIN
-        # data region which is non-contiguous in memory. The performance hit from
-        # dealing with non-contiguous memory is so big (introduces >1 copy), it's
-        # better to checkpoint unneccesarry stuff to get a contiguous chunk of memory.
-        ptr = symbol._data[timestep - i, :, :]
-        ptrs.append(ptr)
-    return ptrs
+from devito.checkpointing import *  # noqa

examples/seismic/acoustic/wavesolver.py

@@ -1,9 +1,8 @@
-from devito import Function, TimeFunction
+from devito import Function, TimeFunction, DevitoCheckpoint, CheckpointOperator
 from devito.tools import memoized_meth
 from examples.seismic.acoustic.operators import (
     ForwardOperator, AdjointOperator, GradientOperator, BornOperator
 )
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
 from pyrevolve import Revolver
 
 

examples/seismic/tti/wavesolver.py

@@ -1,10 +1,9 @@
 # coding: utf-8
-from devito import Function, TimeFunction, warning
+from devito import Function, TimeFunction, warning, DevitoCheckpoint, CheckpointOperator
 from devito.tools import memoized_meth
 from examples.seismic.tti.operators import ForwardOperator, AdjointOperator
 from examples.seismic.tti.operators import JacobianOperator, JacobianAdjOperator
 from examples.seismic.tti.operators import particle_velocity_fields
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
 from pyrevolve import Revolver
 
 

examples/seismic/viscoacoustic/wavesolver.py

@@ -1,10 +1,10 @@
-from devito import VectorTimeFunction, TimeFunction, Function, NODE
+from devito import (VectorTimeFunction, TimeFunction, Function, NODE,
+                    DevitoCheckpoint, CheckpointOperator)
 from devito.tools import memoized_meth
 from examples.seismic import PointSource
 from examples.seismic.viscoacoustic.operators import (
     ForwardOperator, AdjointOperator, GradientOperator, BornOperator
 )
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
 from pyrevolve import Revolver
 
 

tests/test_checkpointing.py

@@ -4,8 +4,8 @@
 from pyrevolve import Revolver
 import numpy as np
 
-from devito import Grid, TimeFunction, Operator, Function, Eq, switchconfig, Constant
-from examples.checkpointing.checkpoint import DevitoCheckpoint, CheckpointOperator
+from devito import (Grid, TimeFunction, Operator, Function, Eq, switchconfig, Constant,
+                    DevitoCheckpoint, CheckpointOperator)
 from examples.seismic.acoustic.acoustic_example import acoustic_setup