compiler: Start adding estimate_memory utility

Author: EdCaunt

devito/operator/operator.py

@@ -7,6 +7,7 @@
 from tempfile import gettempdir
 
 from sympy import sympify
+from sympy import Basic as SympyBasic
 import numpy as np
 
 from devito.arch import ANYCPU, Device, compiler_registry, platform_registry
@@ -40,6 +41,9 @@
 __all__ = ['Operator']
 
 
+_layers = (disk_layer, host_layer, device_layer)
+
+
 class Operator(Callable):
 
     """
@@ -597,6 +601,8 @@ def _prepare_arguments(self, autotune=None, **kwargs):
 
         # Prepare to process data-carriers
         args = kwargs['args'] = ReducerMap()
+
+        # TODO: Add 'estimate-memory' here
         kwargs['metadata'] = {'language': self._language,
                               'platform': self._platform,
                               'transients': self.transients,
@@ -644,9 +650,6 @@ def _prepare_arguments(self, autotune=None, **kwargs):
         for i in discretizations:
             args.update(i._arg_values(**kwargs))
 
-        # TODO: Want to be able to simply stop at this stage and get
-        # the ArgumentsMap for processing
-
         # An ArgumentsMap carries additional metadata that may be used by
         # the subsequent phases of the arguments processing
         args = kwargs['args'] = ArgumentsMap(args, grid, self)
@@ -860,6 +863,38 @@ def cinterface(self, force=False):
     def __call__(self, **kwargs):
         return self.apply(**kwargs)
 
+    def estimate_memory(self, human_readable=True, **kwargs):
+        """
+        Estimate the memory consumed by the Operator.
+
+        TODO: Finish this docstring
+        """
+        # Build the arguments list for which to get the memory consumption
+        # This is so that the estimate will factor in overrides
+        args = self.arguments(**kwargs)
+        mem = args.nbytes_avail_mapper
+
+        if human_readable:
+            # TODO: Fill real values in here
+            # TODO: Format these values to have 3 digits and suitable units
+
+            headline = f"Memory consumption for operator `{self.name}`:"
+            # Table is 28 characters wide
+            lpad = " "*((len(headline) - 28) // 2)
+            info(
+                "\n"
+                f"{headline}\n"
+                f"{lpad}┌────────┬────────┬────────┐\n"
+                f"{lpad}│  Disk  │  Host  │ Device │\n"
+                f"{lpad}├────────┼────────┼────────┤\n"
+                f"{lpad}│   1    │   2    │   3    │\n"
+                f"{lpad}└────────┴────────┴────────┘\n"
+            )
+        else:
+            info(f"{self.name} {mem[disk_layer]} {mem[host_layer]} {mem[device_layer]}")
+
+        from IPython import embed; embed()
+
     def apply(self, **kwargs):
         """
         Execute the Operator.
@@ -1251,6 +1286,7 @@ def nbytes_avail_mapper(self):
         """
         mapper = {}
 
+        # TODO: This doesn't account for the size of the snapshots?
         # The amount of space available on the disk
         usage = shutil.disk_usage(gettempdir())
         mapper[disk_layer] = usage.free
@@ -1267,26 +1303,95 @@ def nbytes_avail_mapper(self):
             nproc = 1
         mapper[host_layer] = int(ANYCPU.memavail() / nproc)
 
-        for layer, consumed in zip((host_layer, device_layer), self.nbytes_consumed):
-            mapper[layer] -= consumed
+        for layer in (host_layer, device_layer):
+            try:
+                mapper[layer] -= self.nbytes_consumed_operator[layer]
+            except KeyError:
+                continue
 
         mapper = {k: int(v) for k, v in mapper.items()}
 
         return mapper
 
     # TODO: This will want some suitable tests in due course
+    # TODO: Might want to also check the spillover onto disk
     @cached_property
     def nbytes_consumed(self):
-        consumed_host, consumed_device = self.nbytes_consumed_heap
-        return consumed_host, consumed_device + self.nbytes_consumed_memmapped
+        """Memory consumed by all objects in the operator"""
+        mem_locations = (
+            self.nbytes_consumed_function,
+            self.nbytes_consumed_array,
+            self.nbytes_consumed_memmapped
+        )
+        return {layer: sum(loc[layer] for loc in mem_locations) for layer in _layers}
+
+    @cached_property
+    def nbytes_consumed_operator(self):
+        """Memory consumed by objects allocated within the operator"""
+        mem_locations = (
+            self.nbytes_consumed_array,
+            self.nbytes_consumed_memmapped
+        )
+        return {layer: sum(loc[layer] for loc in mem_locations) for layer in _layers}
+
+    @cached_property
+    def nbytes_consumed_function(self):
+        """
+        Memory consumed on both device and host by Functions in the
+        corresponding operator.
+        """
+        def get_nbytes(obj):
+            if obj.is_regular:
+                nbytes = obj.nbytes
+            else:
+                nbytes = obj.nbytes_max
+
+            # Could nominally have symbolic nbytes at this point
+            if isinstance(nbytes, SympyBasic):
+                return subs_op_args(nbytes, self)
+            else:
+                return nbytes
+
+        host = 0
+        device = 0
+
+        # Symbols in the operator which may or may not carry data
+        op_symbols = FindSymbols().visit(self.op)
+
+        # Filter out arrays, aliases and non-AbstractFunction objects
+        op_symbols = [i for i in op_symbols if i.is_AbstractFunction
+                      and not i.is_Array and not i.alias]
+
+        for i in op_symbols:
+            # FIXME: Probably wrong for streamed functions
+            # Will overreport memory usage currently
+            try:
+                v = get_nbytes(self[i.name]._obj)
+            except AttributeError:
+                v = get_nbytes(i)
+
+            if i._mem_host:
+                host += v
+            elif i._mem_local:
+                if isinstance(self.platform, Device):
+                    device += v
+                else:
+                    host += v
+            elif i._mem_mapped:
+                if isinstance(self.platform, Device):
+                    device += v
+                host += v
+
+        return {disk_layer: 0, host_layer: host, device_layer: device}
 
     @cached_property
-    def nbytes_consumed_heap(self):
+    def nbytes_consumed_array(self):
         """
-        Memory consumed on both device and host by the corresponding operator.
+        Memory consumed on both device and host by C-land Arrays
+        in the corresponding operator.
         """
-        host_layer = 0
-        device_layer = 0
+        host = 0
+        device = 0
 
         # Temporaries such as Arrays are allocated and deallocated on-the-fly
         # while in C land, so they need to be accounted for as well
@@ -1304,26 +1409,26 @@ def nbytes_consumed_heap(self):
                 continue
 
             if i._mem_host:
-                host_layer += v
+                host += v
             elif i._mem_local:
                 if isinstance(self.platform, Device):
-                    device_layer += v
+                    device += v
                 else:
-                    host_layer += v
+                    host += v
             elif i._mem_mapped:
                 if isinstance(self.platform, Device):
-                    device_layer += v
-                host_layer += v
+                    device += v
+                host += v
 
-        return host_layer, device_layer
+        return {disk_layer: 0, host_layer: host, device_layer: device}
 
     @cached_property
     def nbytes_consumed_memmapped(self):
         """
         Memory also consumed on device by data which is to be memcpy-d
         from host to device at the start of computation.
         """
-        device_layer = 0
+        device = 0
         # All input Functions are yet to be memcpy-ed to the device
         # TODO: this may not be true depending on `devicerm`, which is however
         # virtually never used
@@ -1337,11 +1442,11 @@ def nbytes_consumed_memmapped(self):
                             v = self[i.name]._obj.nbytes
                         except AttributeError:
                             v = i.nbytes
-                        device_layer += v
+                        device += v
                 except AttributeError:
                     pass
 
-        return device_layer
+        return {disk_layer: 0, host_layer: 0, device_layer: device}
 
 
 def parse_kwargs(**kwargs):