compiler: Minimize code size reusing efuncs as much as possible

devito/ir/iet/utils.py

@@ -2,7 +2,7 @@
 
 from devito.ir.iet import FindSections, FindSymbols
 from devito.symbolics import Keyword, Macro
-from devito.tools import filter_ordered
+from devito.tools import filter_ordered, natural_sort_key
 from devito.types import Global
 
 __all__ = [
@@ -39,6 +39,10 @@ def __getitem__(self, key):
         return IterationTree(ret) if isinstance(key, slice) else ret
 
 
+def _canonical_parameter_key(parameter):
+    return str(type(parameter)), natural_sort_key(parameter.name)
+
+
 def retrieve_iteration_tree(node, mode='normal'):
     """
     A list with all Iteration sub-trees within an IET.
@@ -138,7 +142,7 @@ def derive_parameters(iet, drop_locals=False, ordering='default'):
     # amount of tests and examples; plus, it might break compatibility those
     # using devito as a library-generator to be embedded within legacy codes
     if ordering == 'canonical':
-        parameters = sorted(parameters, key=lambda p: str(type(p)))
+        parameters = sorted(parameters, key=_canonical_parameter_key)
 
     return parameters
 

devito/ir/iet/visitors.py

@@ -27,7 +27,7 @@
 from devito.symbolics.extended_dtypes import NoDeclStruct
 from devito.tools import (
     GenericVisitor, as_tuple, c_restrict_void_p, filter_ordered, filter_sorted, flatten,
-    is_external_ctype, memoized_weak_meth, sorted_priority
+    is_external_ctype, memoized_weak_meth, natural_sort_key
 )
 from devito.types import (
     ArrayObject, CompositeObject, DeviceMap, Dimension, IndexedData, Pointer
@@ -1127,7 +1127,7 @@ def visit(self, o, *args, **kwargs):
         return super().visit(o, *args, **kwargs)
 
     def _post_visit(self, ret):
-        return sorted(filter_ordered(ret, key=id), key=str)
+        return sorted(filter_ordered(ret, key=id), key=natural_sort_key)
 
     def visit_Node(self, o: Node) -> Iterator[Any]:
         yield from self._visit(o.children)
@@ -1656,14 +1656,4 @@ def generate(self):
 
 
 def sorted_efuncs(efuncs):
-    from devito.ir.iet.efunc import (
-        CommCallable, DeviceFunction, ElementalFunction, ThreadCallable
-    )
-
-    priority = {
-        DeviceFunction: 3,
-        ThreadCallable: 2,
-        ElementalFunction: 1,
-        CommCallable: 1
-    }
-    return sorted_priority(efuncs, priority)
+    return sorted(efuncs, key=lambda i: natural_sort_key(i.name))

devito/passes/clusters/derivatives.py

@@ -9,7 +9,7 @@
 from devito.symbolics import BasicWrapperMixin, reuse_if_untouched, search, uxreplace
 from devito.symbolics.inspection import sympy_dtype
 from devito.tools import infer_dtype, timed_pass
-from devito.types import Eq, Inc, Indexed, Symbol
+from devito.types import Eq, Inc, Indexed, Symbol, Temp
 
 __all__ = ['lower_index_derivatives']
 
@@ -169,15 +169,15 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
     extra = (ispace.itdims + ispace0.itdims,)
     ispace1 = IterationSpace.union(ispace, ispace0, relations=extra)
 
-    # The Symbol that will hold the result of the IndexDerivative computation
+    # The temporary that will hold the result of the IndexDerivative computation
     # NOTE: created before recursing so that we ultimately get a sound ordering
     dtype = sympy_dtype(ideriv)
     try:
         s = reusables.pop()
         assert np.can_cast(s.dtype, dtype)
     except KeyError:
         name = sregistry.make_name(prefix='r')
-        s = Symbol(name=name, dtype=dtype)
+        s = Temp(name=name, dtype=dtype)
 
     # Go inside `expr` and recursively lower any nested IndexDerivatives
     expr, processed = _core(expr, c, ispace1, weights, reusables, mapper, **kwargs)

devito/passes/iet/engine.py

@@ -16,7 +16,8 @@
 from devito.passes import needs_transfer
 from devito.symbolics import FieldFromComposite, FieldFromPointer, IndexedPointer, search
 from devito.tools import (
-    DAG, as_hashable, as_tuple, filter_ordered, memoized_func, sorted_priority, timed_pass
+    DAG, as_hashable, as_tuple, filter_ordered, memoized_func, natural_sort_key,
+    sorted_priority, timed_pass
 )
 from devito.types import (
     Array, Auto, Bundle, ComponentAccess, CompositeObject, FunctionMap, IncrDimension,
@@ -531,6 +532,7 @@ def abstract_efunc(efunc):
             - Objects are renamed as "o0", "o1", ...
     """
     functions = FindSymbols('basics|symbolics|dimensions').visit(efunc)
+    functions = sorted(functions, key=natural_sort_key)
 
     mapper = abstract_objects(tuple(functions))
 

devito/passes/iet/orchestration.py

@@ -141,7 +141,7 @@ def process(self, iet):
             CollectTasks(task_groups).visit(iet)
 
         # Lower the SyncSpots in a single bottom-up traversal, atomically lowering
-        lowerer = LowerSyncSpots(callbacks, task_groups)
+        lowerer = LowerSyncSpots(callbacks, task_groups, self.sregistry)
         iet = lowerer.visit(iet)
 
         return iet, {'efuncs': lowerer.efuncs}
@@ -217,8 +217,7 @@ def visit_Conditional(self, o, **kwargs):
         kwargs['condition'] = o
         self._visit(o.children, **kwargs)
 
-    def visit_SyncSpot(self, o, iteration=None, condition=None,
-                       in_snapshot=False):
+    def visit_SyncSpot(self, o, iteration=None, condition=None, snapshot=None):
         if iteration is not None:
             syncs = as_mapper(o.sync_ops, type)
 
@@ -240,18 +239,19 @@ def visit_SyncSpot(self, o, iteration=None, condition=None,
                 if gid is not None:
                     task = Task(o, guard, sync_ops)
                     self._task_groups.add(
-                        task, condition or o, iteration, gid, optype, in_snapshot
+                        task, snapshot or condition or o, iteration, gid, optype,
+                        snapshot is not None
                     )
 
                 break
 
         if any(isinstance(i, SnapOut) for i in o.sync_ops):
             # Do not mix composite tasks with other compatible groups
-            in_snapshot = True
+            snapshot = o
 
         self._visit(
             o.children, iteration=iteration, condition=condition,
-            in_snapshot=in_snapshot
+            snapshot=snapshot
         )
 
 
@@ -272,10 +272,11 @@ class LowerSyncSpots(Transformer):
         The task groups to lower atomically.
     """
 
-    def __init__(self, callbacks, task_groups):
+    def __init__(self, callbacks, task_groups, sregistry):
         super().__init__({})
 
         self._callbacks = callbacks
+        self._sregistry = sregistry
 
         self._task_bodies = dict.fromkeys(task_groups.sync_spots)
         self._task_groups = task_groups.by_anchor
@@ -346,6 +347,13 @@ def _lower_task_groups(self, o, iet):
                     List(body=self._task_bodies.pop(task.spot)),
                     task.sync_ops, layer, wrap=False
                 )
+
+                if len(tasks) > 1:
+                    name = self._sregistry.make_name(prefix=f'{prefix}_body')
+                    efunc = make_callable(name, task_body)
+                    self._efuncs.append(efunc)
+                    task_body = Call(name, efunc.parameters)
+
                 if task.guard is None:
                     # Preserve the local scope of an unguarded task body
                     scope = Block(body=task_body)

devito/tools/utils.py

@@ -1,3 +1,4 @@
+import re
 import types
 from collections import OrderedDict
 from collections.abc import Iterable, Mapping
@@ -27,6 +28,7 @@
     'invert',
     'is_integer',
     'is_number',
+    'natural_sort_key',
     'powerset',
     'pprint',
     'prod',
@@ -41,6 +43,16 @@
 ]
 
 
+def natural_sort_key(value):
+    """
+    Return a key that sorts embedded decimal numbers numerically.
+    """
+    return tuple(
+        (part.isdigit(), int(part) if part.isdigit() else part)
+        for part in re.split(r'(\d+)', str(value))
+    )
+
+
 def prod(iterable, initial=1):
     return reduce(mul, iterable, initial)
 

examples/mpi/overview.ipynb

@@ -447,10 +447,10 @@
        "  double section0;\n",
        "} ;\n",
        "\n",
-       "static void sendrecv0(struct dataobj *restrict u_vec, const int x_size, const int y_size, int ogtime, int ogx, int ogy, int ostime, int osx, int osy, int fromrank, int torank, MPI_Comm comm);\n",
-       "static void haloupdate0(struct dataobj *restrict u_vec, MPI_Comm comm, struct neighborhood * nb, int otime);\n",
        "static void gather0(float *restrict buf_vec, int bx_size, int by_size, struct dataobj *restrict u_vec, const int otime, const int ox, const int oy);\n",
+       "static void haloupdate0(struct dataobj *restrict u_vec, MPI_Comm comm, struct neighborhood * nb, int otime);\n",
        "static void scatter0(float *restrict buf_vec, int bx_size, int by_size, struct dataobj *restrict u_vec, const int otime, const int ox, const int oy);\n",
+       "static void sendrecv0(struct dataobj *restrict u_vec, const int x_size, const int y_size, int ogtime, int ogx, int ogy, int ostime, int osx, int osy, int fromrank, int torank, MPI_Comm comm);\n",
        "\n",
        "int Kernel(struct dataobj *restrict u_vec, const float h_x, const int time_M, const int time_m, const int x_M, const int x_m, const int y_M, const int y_m, MPI_Comm comm, struct neighborhood * nb, struct profiler * timers)\n",
        "{\n",
@@ -480,38 +480,6 @@
        "  return 0;\n",
        "}\n",
        "\n",
-       "static void sendrecv0(struct dataobj *restrict u_vec, const int x_size, const int y_size, int ogtime, int ogx, int ogy, int ostime, int osx, int osy, int fromrank, int torank, MPI_Comm comm)\n",
-       "{\n",
-       "  MPI_Request rrecv;\n",
-       "  MPI_Request rsend;\n",
-       "\n",
-       "  float *restrict bufg_vec __attribute__ ((aligned (64)));\n",
-       "  posix_memalign((void**)(&bufg_vec),64,sizeof(float)*(long)y_size*(long)x_size);\n",
-       "  float *restrict bufs_vec __attribute__ ((aligned (64)));\n",
-       "  posix_memalign((void**)(&bufs_vec),64,sizeof(float)*(long)y_size*(long)x_size);\n",
-       "\n",
-       "  MPI_Irecv(bufs_vec,x_size*y_size,MPI_FLOAT,fromrank,13,comm,&rrecv);\n",
-       "  if (torank != MPI_PROC_NULL)\n",
-       "  {\n",
-       "    gather0(bufg_vec,x_size,y_size,u_vec,ogtime,ogx,ogy);\n",
-       "  }\n",
-       "  MPI_Isend(bufg_vec,x_size*y_size,MPI_FLOAT,torank,13,comm,&rsend);\n",
-       "  MPI_Wait(&rsend,MPI_STATUS_IGNORE);\n",
-       "  MPI_Wait(&rrecv,MPI_STATUS_IGNORE);\n",
-       "  if (fromrank != MPI_PROC_NULL)\n",
-       "  {\n",
-       "    scatter0(bufs_vec,x_size,y_size,u_vec,ostime,osx,osy);\n",
-       "  }\n",
-       "\n",
-       "  free(bufg_vec);\n",
-       "  free(bufs_vec);\n",
-       "}\n",
-       "\n",
-       "static void haloupdate0(struct dataobj *restrict u_vec, MPI_Comm comm, struct neighborhood * nb, int otime)\n",
-       "{\n",
-       "  sendrecv0(u_vec,u_vec->hsize[3],u_vec->npsize[2],otime,u_vec->oofs[2],u_vec->hofs[4],otime,u_vec->hofs[3],u_vec->hofs[4],nb->rc,nb->lc,comm);\n",
-       "}\n",
-       "\n",
        "static void gather0(float *restrict buf_vec, int bx_size, int by_size, struct dataobj *restrict u_vec, const int otime, const int ox, const int oy)\n",
        "{\n",
        "  float (*restrict buf)[bx_size][by_size] __attribute__ ((aligned (64))) = (float (*)[bx_size][by_size]) buf_vec;\n",
@@ -532,6 +500,11 @@
        "  }\n",
        "}\n",
        "\n",
+       "static void haloupdate0(struct dataobj *restrict u_vec, MPI_Comm comm, struct neighborhood * nb, int otime)\n",
+       "{\n",
+       "  sendrecv0(u_vec,u_vec->hsize[3],u_vec->npsize[2],otime,u_vec->oofs[2],u_vec->hofs[4],otime,u_vec->hofs[3],u_vec->hofs[4],nb->rc,nb->lc,comm);\n",
+       "}\n",
+       "\n",
        "static void scatter0(float *restrict buf_vec, int bx_size, int by_size, struct dataobj *restrict u_vec, const int otime, const int ox, const int oy)\n",
        "{\n",
        "  float (*restrict buf)[bx_size][by_size] __attribute__ ((aligned (64))) = (float (*)[bx_size][by_size]) buf_vec;\n",
@@ -551,6 +524,33 @@
        "    }\n",
        "  }\n",
        "}\n",
+       "\n",
+       "static void sendrecv0(struct dataobj *restrict u_vec, const int x_size, const int y_size, int ogtime, int ogx, int ogy, int ostime, int osx, int osy, int fromrank, int torank, MPI_Comm comm)\n",
+       "{\n",
+       "  MPI_Request rrecv;\n",
+       "  MPI_Request rsend;\n",
+       "\n",
+       "  float *restrict bufg_vec __attribute__ ((aligned (64)));\n",
+       "  posix_memalign((void**)(&bufg_vec),64,sizeof(float)*(long)y_size*(long)x_size);\n",
+       "  float *restrict bufs_vec __attribute__ ((aligned (64)));\n",
+       "  posix_memalign((void**)(&bufs_vec),64,sizeof(float)*(long)y_size*(long)x_size);\n",
+       "\n",
+       "  MPI_Irecv(bufs_vec,x_size*y_size,MPI_FLOAT,fromrank,13,comm,&rrecv);\n",
+       "  if (torank != MPI_PROC_NULL)\n",
+       "  {\n",
+       "    gather0(bufg_vec,x_size,y_size,u_vec,ogtime,ogx,ogy);\n",
+       "  }\n",
+       "  MPI_Isend(bufg_vec,x_size*y_size,MPI_FLOAT,torank,13,comm,&rsend);\n",
+       "  MPI_Wait(&rsend,MPI_STATUS_IGNORE);\n",
+       "  MPI_Wait(&rrecv,MPI_STATUS_IGNORE);\n",
+       "  if (fromrank != MPI_PROC_NULL)\n",
+       "  {\n",
+       "    scatter0(bufs_vec,x_size,y_size,u_vec,ostime,osx,osy);\n",
+       "  }\n",
+       "\n",
+       "  free(bufg_vec);\n",
+       "  free(bufs_vec);\n",
+       "}\n",
        "\n"
       ]
      },

tests/test_iet.py

@@ -11,18 +11,20 @@
 from devito.ir import SymbolRegistry
 from devito.ir.iet import (
     Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr,
-    ElementalFunction, FindSymbols, Iteration, KernelLaunch, Lambda, List, Switch,
-    Transformer, filter_iterations, make_efunc, retrieve_iteration_tree
+    ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List,
+    Switch, Transformer, filter_iterations, make_callable, make_efunc,
+    retrieve_iteration_tree
 )
-from devito.passes.iet.engine import Graph
+from devito.ir.iet.visitors import sorted_efuncs
+from devito.passes.iet.engine import Graph, reuse_efuncs
 from devito.passes.iet.languages.C import CDataManager
 from devito.symbolics import (
     FLOAT, Byref, Class, FieldFromComposite, InlineIf, ListInitializer, Macro, SizeOf,
     String
 )
 from devito.symbolics.extended_dtypes import c_complex
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
-from devito.types import Array, CustomDimension, LocalObject, Pointer, Symbol
+from devito.types import Array, CustomDimension, LocalObject, Pointer, Symbol, Temp
 from devito.types.misc import FunctionMap
 
 
@@ -134,6 +136,15 @@ def test_find_symbols_nested(mode, expected):
     assert [f.name for f in found] == eval(expected)
 
 
+def test_find_symbols_natural_numbering():
+    temporaries = [Temp(name=f'r{i}') for i in range(12)]
+    call = Call('foo', temporaries[::-1])
+
+    found = FindSymbols('basics').visit(call)
+
+    assert [i.name for i in found] == [i.name for i in temporaries]
+
+
 def test_list_denesting():
     l0 = List(header=cgen.Line('a'), body=List(header=cgen.Line('b')))
     l1 = l0._rebuild(body=List(header=cgen.Line('c')))
@@ -529,6 +540,37 @@ def test_codegen_quality0():
     assert foo1.parameters[0] is a
 
 
+def test_reuse_efuncs_natural_numbering():
+
+    def make_foo(name, start):
+        pointers = [Pointer(name=f'p{i}') for i in range(start, start + 3)]
+        return make_callable(name, Call('bar', pointers))
+
+    foo0 = make_foo('foo0', 0)
+    foo1 = make_foo('foo1', 9)
+
+    root = make_callable('root', List(body=[
+        Call(foo0.name, foo0.parameters),
+        Call(foo1.name, foo1.parameters)
+    ]))
+    efuncs = {i.name: i for i in (root, foo0, foo1)}
+
+    efuncs = reuse_efuncs(root, efuncs)
+
+    assert set(efuncs) == {'root', 'foo0'}
+    calls = FindNodes(Call).visit(efuncs[root.name])
+    assert [i.name for i in calls] == ['foo0', 'foo0']
+
+
+def test_sorted_efuncs_natural_numbering():
+    names = ['kernel3', 'foo0', 'kernel11', 'kernel1', 'kernel2']
+    efuncs = [make_callable(name, List()) for name in names]
+
+    assert [i.name for i in sorted_efuncs(efuncs)] == [
+        'foo0', 'kernel1', 'kernel2', 'kernel3', 'kernel11'
+    ]
+
+
 def test_complex_array():
     grid = Grid(shape=(4, 4, 4))
     _, y, z = grid.dimensions