compiler: Preserve compilation invariants under symbol renaming

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
 

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]

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')

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)
 

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)

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

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)
 

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)
 
 

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",

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

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),

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: