diff --git a/devito/ir/cgen/printer.py b/devito/ir/cgen/printer.py index 27ee0e3112..adfc7651aa 100644 --- a/devito/ir/cgen/printer.py +++ b/devito/ir/cgen/printer.py @@ -289,7 +289,7 @@ def _print_Mul(self, expr): term = self.parenthesize(args[0], precedence(expr)) if neg: - return f'-{term}' + return term[1:] if term.startswith('-') else f'-{term}' else: return term diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 6186ff4e76..8a825c7619 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -1506,7 +1506,8 @@ def visit_PointerCast(self, o): def visit_Dereference(self, o): pointee = self.mapper.get(o.pointee, o.pointee) pointer = self.mapper.get(o.pointer, o.pointer) - return reuse_if_unchanged(o, pointee=pointee, pointer=pointer) + offset = uxreplace(o.offset, self.mapper) + return reuse_if_unchanged(o, pointee=pointee, pointer=pointer, offset=offset) def visit_Pragma(self, o): arguments = [uxreplace(i, self.mapper) for i in o.arguments] diff --git a/devito/ir/support/symregistry.py b/devito/ir/support/symregistry.py index d7fca20f7c..01ceaf3bdf 100644 --- a/devito/ir/support/symregistry.py +++ b/devito/ir/support/symregistry.py @@ -1,4 +1,5 @@ -from devito.tools import generator +from devito.symbolics import retrieve_functions +from devito.tools import as_tuple, generator from devito.types import ( DeviceID, NPThreads, NThreads, NThreadsNested, NThreadsNonaffine, ThreadID ) @@ -15,6 +16,7 @@ class SymbolRegistry: def __init__(self): # {name -> generator()} -- to create unique names for symbols, functions, ... self.counters = {} + self.reserved = set() # Special symbols self.nthreads = NThreads() @@ -31,6 +33,14 @@ def __init__(self): # passes, to maximize symbol (especially Dimension) reuse self.caches = {} + def reserve(self, exprs): + """Reserve identifiers already used by the input expressions.""" + exprs = as_tuple(exprs) + self.reserved.update(i.name for e in exprs for i in e.free_symbols) + + # Before indexification, a Function's free symbols omit its own name + self.reserved.update(i.name for i in retrieve_functions(exprs)) + def make_name(self, prefix=None, increment_first=True): # By default we're creating a new symbol if prefix is None: @@ -42,10 +52,15 @@ def make_name(self, prefix=None, increment_first=True): counter = self.counters.setdefault(prefix, generator()) # Only increment symbol names after the first encountered - if not increment_first: + if not increment_first and prefix not in self.reserved: + self.reserved.add(prefix) return prefix - return f'{prefix}{counter()}' + name = f'{prefix}{counter()}' + while name in self.reserved: + name = f'{prefix}{counter()}' + self.reserved.add(name) + return name def make_npthreads(self, size): name = self.make_name(prefix='npthreads') diff --git a/devito/operator/operator.py b/devito/operator/operator.py index 8f3c02f983..e36c48a351 100644 --- a/devito/operator/operator.py +++ b/devito/operator/operator.py @@ -293,6 +293,8 @@ def _lower(cls, expressions, **kwargs): expressions = as_tuple(expressions) + kwargs['sregistry'].reserve(expressions) + # Enable recursive lowering # This may be used by a compilation pass that constructs a new # expression for which a partial or complete lowering is desired @@ -472,6 +474,11 @@ def _specialize_iet(cls, graph, **kwargs): """ return graph + @classmethod + def _finalize_iet(cls, graph, **kwargs): + """Backend hook after lowering, before macro and signature extraction.""" + return graph + @classmethod @timed_pass(name='lowering.uIET') def _lower_uiet(cls, stree, profiler=None, **kwargs): @@ -515,15 +522,17 @@ def _lower_iet(cls, uiet, **kwargs): # specialization further Sections may be introduced cls._Target.instrument(graph, **kwargs) - # Extract the necessary macros from the symbolic objects - generate_macros(graph, **kwargs) - # Target-specific lowering lower_dtypes(graph, **kwargs) # Target-independent optimizations minimize_symbols(graph) + graph = cls._finalize_iet(graph, **kwargs) + + # Extract macros only after the symbolic objects have their final names + generate_macros(graph, **kwargs) + # Finalize helper signatures after all IET transformations have settled. finalize_args(graph) diff --git a/devito/passes/clusters/cse.py b/devito/passes/clusters/cse.py index 9c899eb23e..75fe68d636 100644 --- a/devito/passes/clusters/cse.py +++ b/devito/passes/clusters/cse.py @@ -293,15 +293,13 @@ def _toposort(exprs): if e0.lhs in e1.rhs.free_symbols: dag.add_edge(e0, e1, force_add=True) + ordering = {e: n for n, e in enumerate(exprs)} + def choose_element(queue, scheduled): + # Preserve expression order among ready temporaries, independent of names tmps = [i for i in queue if isinstance(i.lhs, CTemp)] - if tmps: - # Try to honor temporary names as much as possible - first = sorted(tmps, key=lambda i: i.lhs.name).pop(0) - queue.remove(first) - else: - first = sorted(queue, key=lambda i: exprs.index(i)).pop(0) - queue.remove(first) + first = min(tmps or queue, key=ordering.__getitem__) + queue.remove(first) return first processed = dag.topological_sort(choose_element) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index fe9a78b65b..8615582da5 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -564,6 +564,8 @@ def abstract_objects(objects0): mapper = {} sregistry = SymbolRegistry() for i in objects: + if i in mapper: + continue abstract_object(i, mapper, sregistry) return mapper diff --git a/devito/passes/iet/misc.py b/devito/passes/iet/misc.py index 1b404daca1..a94a754124 100644 --- a/devito/passes/iet/misc.py +++ b/devito/passes/iet/misc.py @@ -188,8 +188,9 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): subs = {} for i in findexeds: + key = (i.base, i.strides_map) try: - v = tracker[i.base].v + v = tracker[key].v subs[i] = v.func(v.base, *i.indices) continue except KeyError: @@ -199,7 +200,7 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs): header, v = i.bind(pname) subs[i] = v - tracker[i.base] = Bunch(header=header, v=v) + tracker[key] = Bunch(header=header, v=v) iet = Uxreplace(subs).visit(iet) diff --git a/devito/symbolics/manipulation.py b/devito/symbolics/manipulation.py index b545290a15..82fffcfb9a 100644 --- a/devito/symbolics/manipulation.py +++ b/devito/symbolics/manipulation.py @@ -12,7 +12,7 @@ EvalDerivative, IndexDerivative, LocalSum ) from devito.symbolics.extended_dtypes import LONG -from devito.symbolics.extended_sympy import DefFunction, rfunc +from devito.symbolics.extended_sympy import CallFromPointer, DefFunction, rfunc from devito.symbolics.queries import q_leaf from devito.symbolics.search import ( retrieve_functions, retrieve_indexed, retrieve_terminals @@ -22,11 +22,12 @@ from devito.symbolics.unevaluation import Pow as UnevalPow from devito.symbolics.unevaluation import UnevaluableMixin from devito.tools import ( - EnrichedTuple, as_list, as_tuple, flatten, split, transitive_closure + EnrichedTuple, as_list, as_tuple, flatten, frozendict, split, transitive_closure ) from devito.types.array import ComponentAccess from devito.types.basic import Basic, Indexed from devito.types.equation import Eq +from devito.types.misc import FIndexed from devito.types.relational import Ge, Gt, Le, Lt __all__ = [ @@ -53,8 +54,9 @@ def uxreplace(expr, rule): be avoided at all costs (e.g., to prevent SymPy from unpicking Devito transformations, such as factorization). - The canonical ordering of the arguments is however guaranteed; where this is - not possible, a re-evaluation will be enforced. + Evaluated expressions retain canonical argument ordering. Lowered, + unevaluated arithmetic retains its operand order and grouping unless + substitution makes it completely numeric. By avoiding re-evaluations, this function is typically much quicker than SymPy's xreplace. @@ -117,6 +119,17 @@ def _uxreplace(expr, rule): return expr, False +@_uxreplace.register(CallFromPointer) +def _(expr, rule): + if expr in rule or expr.__class__ in rule: + return _uxreplace.dispatch(object)(expr, rule) + + args, changed = _uxreplace_dispatch((expr.pointer, expr.params), rule) + if changed: + return expr.func(expr.call, *args), True + return expr, False + + @_uxreplace.register(LocalSum) def _(expr, rule): # Only the sum owns these guards; ordinary ConditionalDimensions stay atomic @@ -165,6 +178,7 @@ def _(iterable, rule): @_uxreplace_dispatch.register(dict) +@_uxreplace_dispatch.register(frozendict) def _(mapper, rule): ret = {} changed = False @@ -183,6 +197,12 @@ def _uxreplace_handle(expr, args, kwargs): return expr.func(*args) +@_uxreplace_handle.register(CallFromPointer) +def _(expr, args, kwargs): + # C member names belong to their compound type, rather than the local scope + return expr.func(expr.call, *args[1:]) + + @_uxreplace_handle.register(Min) @_uxreplace_handle.register(Max) @_uxreplace_handle.register(Pow) @@ -191,6 +211,15 @@ def _(expr, args, kwargs): return expr.func(*args, evaluate=evaluate) +@_uxreplace_handle.register(UnevalAdd) +@_uxreplace_handle.register(UnevalMul) +def _(expr, args, kwargs): + # Lowered arithmetic must survive substitutions without reassociation + if all(i.is_Number for i in args): + return expr.func.__base__(*args) + return expr.func(*args) + + @_uxreplace_handle.register(Add) def _(expr, args, kwargs): args = [i for i in args if i != 0] @@ -252,6 +281,7 @@ def dispatchable(self, obj): _uxreplace_registry.register(Eq) _uxreplace_registry.register(DefFunction) _uxreplace_registry.register(ComponentAccess) +_uxreplace_registry.register(FIndexed) _uxreplace_registry.register(LocalSum) diff --git a/examples/performance/01_gpu.ipynb b/examples/performance/01_gpu.ipynb index dc096aa0eb..d37e956cd6 100644 --- a/examples/performance/01_gpu.ipynb +++ b/examples/performance/01_gpu.ipynb @@ -309,7 +309,7 @@ " {\n", " for (int y = y_m; y <= y_M; y += 1)\n", " {\n", - " uL0(time + 1, x + 2, y + 2) = dt*(c*(r1*uL0(time, x + 1, y + 2) + r1*uL0(time, x + 3, y + 2) + r2*uL0(time, x + 2, y + 1) + r2*uL0(time, x + 2, y + 3) - 2.0F*(r1*uL0(time, x + 2, y + 2) + r2*uL0(time, x + 2, y + 2))) + r0*uL0(time, x + 2, y + 2));\n", + " uL0(time + 1, x + 2, y + 2) = dt*(c*(r1*uL0(time, x + 1, y + 2) + r1*uL0(time, x + 3, y + 2) + r2*uL0(time, x + 2, y + 1) + r2*uL0(time, x + 2, y + 3) + (-2.0F)*(r1*uL0(time, x + 2, y + 2) + r2*uL0(time, x + 2, y + 2))) + r0*uL0(time, x + 2, y + 2));\n", " }\n", " }\n", " STOP(section0,timers)\n", diff --git a/tests/test_cse.py b/tests/test_cse.py index 6ab68f05a0..4fc439c49b 100644 --- a/tests/test_cse.py +++ b/tests/test_cse.py @@ -8,7 +8,7 @@ ConditionalDimension, Eq, Function, Grid, Operator, TimeFunction, cos, sin ) from devito.finite_differences.differentiable import diffify -from devito.ir import Conditional, DummyEq, FindNodes, FindSymbols +from devito.ir import Conditional, DummyEq, Expression, FindNodes, FindSymbols from devito.ir.support import generator from devito.passes.clusters.cse import CTemp, _cse from devito.symbolics import DefFunction, FieldFromPointer, IndexedPointer, indexify @@ -131,6 +131,23 @@ def test_temp_order(): assert type(args[2]) is CTemp +@pytest.mark.parametrize('names', [('r2', 'r3'), ('r9', 'r10'), ('later', 'earlier')]) +@pytest.mark.parametrize('mode', ['smartsort', 'advanced']) +def test_schedule_independent_of_temp_names(names, mode): + grid = Grid(shape=(4, 4)) + f = Function(name='f', grid=grid) + g = Function(name='g', grid=grid) + u = Function(name='u', grid=grid) + v = Function(name='v', grid=grid) + a, b = [CTemp(name=name, dtype=np.float32) for name in names] + + op = Operator([Eq(a, cos(f)), Eq(b, sin(g)), Eq(u, a + b), Eq(v, a*b)], + opt=('advanced', {'cse-algo': mode})) + + exprs = [e.expr for e in FindNodes(Expression).visit(op)] + assert [e.lhs for e in exprs if isinstance(e.lhs, CTemp)][:2] == [a, b] + + def test_deffunction_not_captured(): """ Opaque C-level calls (DefFunction) must not be CSE-captured: their return diff --git a/tests/test_iet.py b/tests/test_iet.py index 99a94ed24c..2834ac2712 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -12,7 +12,7 @@ from devito.ir.iet import ( Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr, ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List, - Switch, Transformer, filter_iterations, make_callable, make_efunc, + Switch, Transformer, Uxreplace, filter_iterations, make_callable, make_efunc, retrieve_iteration_tree ) from devito.ir.iet.visitors import sorted_efuncs @@ -120,6 +120,20 @@ def test_nested_calls_cgen(): assert str(code) == 'foo(bar());' +def test_uxreplace_dereference_offset(): + d = CustomDimension(name='d', symbolic_size=4) + array = Array(name='a0', dimensions=(d,), shape=(4,), scope='shared') + pointer = Pointer(name='sm7') + old, new = Symbol(name='off7'), Symbol(name='off0') + node = Dereference(array, pointer, offset=old + 32) + + replaced = Uxreplace({old: new}).visit(node) + + assert replaced.offset == new + 32 + assert replaced.pointee is array + assert replaced.pointer is pointer + + @pytest.mark.parametrize('mode,expected', [ ('basics', '["x"]'), ('symbolics', '["f"]') @@ -542,14 +556,17 @@ def test_codegen_quality0(): assert foo1.parameters[0] is a -def test_reuse_efuncs_natural_numbering(): +@pytest.mark.parametrize('copy_symbols', [False, True]) +def test_reuse_efuncs_natural_numbering(copy_symbols): - def make_foo(name, start): + def make_foo(name, start, copies=False): pointers = [Pointer(name=f'p{i}') for i in range(start, start + 3)] - return make_callable(name, Call('bar', pointers)) + aliases = [p._rebuild() for p in pointers] if copies else pointers + return make_callable(name, List(body=[Call('bar', pointers), + Call('bar', aliases)])) foo0 = make_foo('foo0', 0) - foo1 = make_foo('foo1', 9) + foo1 = make_foo('foo1', 9, copies=copy_symbols) root = make_callable('root', List(body=[ Call(foo0.name, foo0.parameters), diff --git a/tests/test_operator.py b/tests/test_operator.py index e6fc86463a..8b34a93bc9 100644 --- a/tests/test_operator.py +++ b/tests/test_operator.py @@ -198,6 +198,19 @@ def test_jit_compiler_flags(self, envvar, flags, monkeypatch): assert op0._soname != op1._soname assert op1._soname == op2._soname + @pytest.mark.parametrize('name', ['r0', 'r1']) + def test_generated_names_avoid_inputs(self, name): + grid = Grid(shape=(8, 8), extent=(7., 7.)) + f = Function(name=name, grid=grid, space_order=4) + u = Function(name='u', grid=grid) + x, y = np.indices(f.data_with_halo.shape, dtype=np.float32) + f.data_with_halo[:] = x + 2*y + + op = Operator(Eq(u, f.dx + f.dy), name='names_inputs') + op.apply() + + np.testing.assert_allclose(u.data, 3., rtol=1.e-5, atol=1.e-6) + class TestCodeGen: diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 51a759a705..23c74b5525 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -15,17 +15,20 @@ from devito.ir.support.guards import GuardExpr, pairwise_or, simplify_and from devito.mpi.halo_scheme import HaloTouch from devito.symbolics import ( # noqa - INT, BaseCast, CallFromPointer, Cast, DefFunction, FieldFromComposite, - FieldFromPointer, IntDiv, ListInitializer, Namespace, ReservedWord, RoundUp, Rvalue, - SizeOf, VectorAccess, evalrel, pow_to_mul, q_negative, q_positive, - retrieve_derivatives, retrieve_functions, retrieve_indexed, subs_if_composite, - uxreplace, xreplace_indices + INT, BaseCast, CallFromComposite, CallFromPointer, Cast, DefFunction, + FieldFromComposite, FieldFromPointer, IntDiv, ListInitializer, Namespace, + ReservedWord, RoundUp, Rvalue, SizeOf, VectorAccess, evalrel, pow_to_mul, q_negative, + q_positive, retrieve_derivatives, retrieve_functions, retrieve_indexed, + subs_if_composite, uxreplace, xreplace_indices ) +from devito.symbolics.unevaluation import Add as UnevalAdd +from devito.symbolics.unevaluation import Mul as UnevalMul from devito.tools import CustomDtype, as_tuple, dtypes_vector_mapper from devito.types import ( Array, Bundle, ComponentAccess, FIndexed, LocalObject, Object, StencilDimension ) from devito.types import Symbol as dSymbol +from devito.types import Temp from devito.types.basic import AbstractSymbol @@ -1123,6 +1126,44 @@ def test_halo_touch(self): assert ht1.args == (g[x, y],) + @pytest.mark.parametrize('op', [UnevalAdd, UnevalMul]) + def test_lowered_arithmetic_order(self, op): + a, b, c = [Temp(name=f'r{i}') for i in (9, 10, 11)] + x, y, z = [Temp(name=f'r{i}') for i in (0, 1, 2)] + expr = op(c, op(b, a)) + + rebuilt = uxreplace(expr, {a: x, b: y, c: z}) + + assert rebuilt.args == (z, op(y, x)) + + @pytest.mark.parametrize('cls,converted', [(CallFromPointer, CallFromComposite), + (CallFromComposite, CallFromPointer), + (FieldFromPointer, FieldFromComposite), + (FieldFromComposite, FieldFromPointer)]) + def test_bound_member_names(self, cls, converted): + a, b = [dSymbol(name=f'r{i}') for i in (9, 0)] + p, q = [dSymbol(name=f'p{i}') for i in (9, 0)] + expr = cls(a, p, [a]) + + assert uxreplace(expr, {a: b, p: q}) == cls(a, q, [b]) + assert uxreplace(expr, {expr: cls(b, q, [b])}) == cls(b, q, [b]) + assert uxreplace(expr, {cls: converted, a: b, p: q}) == converted(a, q, [b]) + + bound = cls(a, p) + assert uxreplace(bound, {a: b}) is bound + + def test_findexed_strides(self): + grid = Grid(shape=(4, 4)) + x, y = grid.dimensions + f = Function(name='f', grid=grid) + a, b = [dSymbol(name=f'y_stride{i}') for i in (9, 0)] + expr = FIndexed(f.indexed, x, y, strides_map={y: a}) + + rebuilt = uxreplace(expr, {a: b}) + + assert rebuilt.strides == (b,) + assert expr.strides == (a,) + def test_minmax(): grid = Grid(shape=(5, 5))