compiler: Improve rcompile

Author: FabioLuporini

conftest.py

@@ -290,8 +290,8 @@ def pytest_runtest_makereport(item, call):
 
 def pytest_make_parametrize_id(config, val, argname):
     """
-    Prevents pytest to make obsucre parameter names (param0, param1, ...)
-    and default to str(val) instead for better log readability.
+    Prevents pytest from making obscure parameter names (param0, param1, ...)
+    and default to sympy.sstr(val) instead for better log readability.
     """
     # First see if it has a name
     if hasattr(val, '__name__'):

devito/operator/operator.py

@@ -18,8 +18,8 @@
                            switch_log_level)
 from devito.ir.equations import LoweredEq, lower_exprs, concretize_subdims
 from devito.ir.clusters import ClusterGroup, clusterize
-from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols,
-                           MetaCall, derive_parameters, iet_build)
+from devito.ir.iet import (Callable, CInterface, EntryFunction, DeviceFunction,
+                           FindSymbols, MetaCall, derive_parameters, iet_build)
 from devito.ir.support import AccessMode, SymbolRegistry
 from devito.ir.stree import stree_build
 from devito.operator.profiling import create_profile
@@ -1270,9 +1270,12 @@ def rcompile(expressions, kwargs, options, target=None):
 
     # Recursive compilation is expensive, so we cache the result because sometimes
     # it is called multiple times for the same input
-    compiled = RCompiles(expressions, cls).compile(**kwargs)
+    irs, byproduct0 = RCompiles(expressions, cls).compile(**kwargs)
 
-    return compiled
+    key = lambda i: isinstance(i, (EntryFunction, DeviceFunction))
+    byproduct = byproduct0.filter(key)
+
+    return irs, byproduct
 
 
 # *** Misc helpers

devito/passes/iet/definitions.py

@@ -91,7 +91,6 @@ def __init__(self, rcompile=None, sregistry=None, platform=None,
         self.rcompile = rcompile
         self.sregistry = sregistry
         self.platform = platform
-        self.alloc_mapped = (options or {}).get('alloc_mapped', True)
 
     def _alloc_object_on_low_lat_mem(self, site, obj, storage):
         """
@@ -178,10 +177,6 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
         Allocate a mapped Array in the host high bandwidth memory.
         """
-        if not self.alloc_mapped:
-            # Mapped array assumed to be preallocated, likely from rcompile
-            return
-
         decl = Definition(obj)
 
         sizeof_dtypeN = SizeOf(obj.indexed._C_typedata)

devito/passes/iet/engine.py

@@ -1,13 +1,13 @@
-from collections import OrderedDict, defaultdict
+from collections import defaultdict
 from functools import partial, singledispatch, wraps
 
 import numpy as np
 from sympy import Mul
 
 from devito.ir.iet import (
     Call, ExprStmt, Expression, Iteration, SyncSpot, AsyncCallable, FindNodes,
-    FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction,
-    ThreadCallable, Uxreplace, derive_parameters
+    FindSymbols, MapNodes, MetaCall, Transformer, EntryFunction, ThreadCallable,
+    Uxreplace, derive_parameters
 )
 from devito.ir.support import SymbolRegistry
 from devito.mpi.distributed import MPINeighborhood
@@ -28,7 +28,31 @@
 __all__ = ['Graph', 'iet_pass', 'iet_visit']
 
 
-class Graph:
+class Byproduct:
+
+    """
+    A Byproduct is a mutable collection of metadata produced by one or more
+    compiler passes.
+
+    This metadata may be used internally or by the caller itself, typically
+    for code generation purposes.
+    """
+
+    def __init__(self, efuncs=None, includes=None, headers=None, namespaces=None,
+                 globs=None):
+        self.efuncs = efuncs or {}
+        self.includes = includes or []
+        self.headers = headers or []
+        self.namespaces = namespaces or []
+        self.globals = globs or []
+
+    @property
+    def funcs(self):
+        return tuple(MetaCall(v, True) for v in self.efuncs.values()
+                     if not isinstance(v, EntryFunction))
+
+
+class Graph(Byproduct):
 
     """
     DAG representation of a call graph.
@@ -49,16 +73,9 @@ class Graph:
     """
 
     def __init__(self, iet, options=None, sregistry=None, **kwargs):
-        self.efuncs = OrderedDict([(iet.name, iet)])
-
         self.sregistry = sregistry
 
-        self.includes = []
-        self.headers = []
-        self.namespaces = []
-        self.globals = []
-
-        # Stash immutable information useful for one or more compiler passes
+        super().__init__({iet.name: iet})
 
         # All written user-level objects
         writes = FindSymbols('writes').visit(iet)
@@ -79,10 +96,6 @@ def __init__(self, iet, options=None, sregistry=None, **kwargs):
     def root(self):
         return self.efuncs[list(self.efuncs).pop(0)]
 
-    @property
-    def funcs(self):
-        return tuple(MetaCall(v, True) for v in self.efuncs.values())[1:]
-
     @property
     def sync_mapper(self):
         """
@@ -146,7 +159,7 @@ def apply(self, func, **kwargs):
             new_efuncs = metadata.get('efuncs', [])
 
             efuncs[i] = efunc
-            efuncs.update(OrderedDict([(i.name, i) for i in new_efuncs]))
+            efuncs.update(dict([(i.name, i) for i in new_efuncs]))
 
             # Update the parameters / arguments lists since `func` may have
             # introduced or removed objects
@@ -176,10 +189,24 @@ def visit(self, func, **kwargs):
         dag = create_call_graph(self.root.name, self.efuncs)
         toposort = dag.topological_sort()
 
-        mapper = OrderedDict([(i, func(self.efuncs[i], **kwargs)) for i in toposort])
+        mapper = dict([(i, func(self.efuncs[i], **kwargs)) for i in toposort])
 
         return mapper
 
+    def filter(self, key):
+        """
+        Return a Byproduct containing only the Callables in the Graph
+        for which `key` evaluates to True. The resulting object cannot be
+        further modified by an IET pass.
+        """
+        return Byproduct(
+            efuncs={i: v for i, v in self.efuncs.items() if key(v)},
+            includes=as_tuple(self.includes),
+            headers=as_tuple(self.headers),
+            namespaces=as_tuple(self.namespaces),
+            globs=as_tuple(self.globals)
+        )
+
 
 def iet_pass(func):
     if isinstance(func, tuple):
@@ -732,7 +759,7 @@ def _filter(v, efunc=None):
 
         return processed
 
-    efuncs = OrderedDict(efuncs)
+    efuncs = dict(efuncs)
     efuncs[root.name] = root._rebuild(parameters=_filter(root.parameters, root))
 
     # Update all call sites to use the new signature