diff --git a/devito/ir/iet/utils.py b/devito/ir/iet/utils.py index 86b23adbbf..5da6ae075c 100644 --- a/devito/ir/iet/utils.py +++ b/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 diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 4208c4b237..81b12fa53b 100644 --- a/devito/ir/iet/visitors.py +++ b/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)) diff --git a/devito/passes/clusters/derivatives.py b/devito/passes/clusters/derivatives.py index ded4a8661e..fafffcedd5 100644 --- a/devito/passes/clusters/derivatives.py +++ b/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,7 +169,7 @@ 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: @@ -177,7 +177,7 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs): 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) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 37334cea9b..fe9a78b65b 100644 --- a/devito/passes/iet/engine.py +++ b/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)) diff --git a/devito/passes/iet/orchestration.py b/devito/passes/iet/orchestration.py index 4f18c7220a..ae2d55766c 100644 --- a/devito/passes/iet/orchestration.py +++ b/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) diff --git a/devito/tools/utils.py b/devito/tools/utils.py index 470be7e79e..0baa8d1f24 100644 --- a/devito/tools/utils.py +++ b/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) diff --git a/examples/mpi/overview.ipynb b/examples/mpi/overview.ipynb index b46f49aaf4..8c209d3bc6 100644 --- a/examples/mpi/overview.ipynb +++ b/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" ] }, diff --git a/tests/test_iet.py b/tests/test_iet.py index e8e8f8444f..6d2b3126c6 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -11,10 +11,12 @@ 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, @@ -22,7 +24,7 @@ ) 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 diff --git a/tests/test_unexpansion.py b/tests/test_unexpansion.py index 026effcf79..dfbb138cb0 100644 --- a/tests/test_unexpansion.py +++ b/tests/test_unexpansion.py @@ -10,11 +10,26 @@ from devito.finite_differences import Weights from devito.ir import Expression, FindNodes, FindSymbols from devito.parameters import configuration, switchconfig -from devito.types import Dimension, Symbol +from devito.types import Dimension, Symbol, Temp class TestLoopScheduling: + def test_index_derivative_temp(self): + grid = Grid(shape=(10, 10)) + + u = Function(name='u', grid=grid, space_order=4) + v = Function(name='v', grid=grid) + + op = Operator(Eq(v, u.dx2), + opt=('advanced', {'expand': False})) + + temporaries = [ + i for i in FindSymbols().visit(op) if i.name.startswith('r') + ] + assert len(temporaries) == 1 + assert isinstance(temporaries[0], Temp) + def test_backward_dt2(self): grid = Grid(shape=(4, 4))