compiler: Preserve compilation invariants under symbol renaming

devito/ir/cgen/printer.py

@@ -280,16 +280,28 @@ def _print_Mod(self, expr):
         return '%'.join(args)
 
     def _print_Mul(self, expr):
-        args = [a for a in expr.args if a != -1]
-        neg = (len(expr.args) - len(args)) % 2
+        args, neg = [], False
+        # SymPy's CodePrinter extracts negative coefficients with the standalone
+        # _keep_coeff helper, which can rebuild a plain SymPy Mul. Since this
+        # is not a method, UnevalMul cannot override it. Extract signs before
+        # delegating to preserve the product's class without patching SymPy or
+        # copying its multiplication printer.
+        for a in expr.args:
+            if a.is_Number and a.is_negative:
+                a = -a
+                neg = not neg
+                if a == 1:
+                    continue
+            args.append(a)
+        args = args or [S.One]
 
         if len(args) > 1:
             term = super()._print_Mul(expr.func(*args, evaluate=False))
         else:
             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, mode='unique'))
+
     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,14 @@ def _toposort(exprs):
             if e0.lhs in e1.rhs.free_symbols:
                 dag.add_edge(e0, e1, force_add=True)
 
+    # Cache positions to avoid repeated linear scans while choosing ready nodes
+    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

@@ -3,7 +3,7 @@
 from functools import singledispatch
 
 import numpy as np
-from sympy import Add, Max, Min, Mul, Pow, S, SympifyError, Tuple, sympify
+from sympy import Add, Max, Min, Mul, Pow, S, SympifyError, Tuple, ceiling, floor, sympify
 from sympy.core import Basic as SympyBasic
 from sympy.core.add import _addsort
 from sympy.core.mul import _mulsort
@@ -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
@@ -21,12 +21,15 @@
 from devito.symbolics.unevaluation import Mul as UnevalMul
 from devito.symbolics.unevaluation import Pow as UnevalPow
 from devito.symbolics.unevaluation import UnevaluableMixin
+from devito.symbolics.unevaluation import ceiling as UnevalCeil
+from devito.symbolics.unevaluation import floor as UnevalFloor
 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 +56,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.
@@ -77,6 +81,10 @@ def uxreplace(expr, rule):
 
 @singledispatch
 def _uxreplace(expr, rule):
+    return _uxreplace_default(expr, rule)
+
+
+def _uxreplace_default(expr, rule):
     if expr in rule:
         v = rule[expr]
         if not isinstance(v, dict):
@@ -117,6 +125,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_default(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
@@ -128,7 +147,7 @@ def _(expr, rule):
             mapper[d] = d._rebuild(**kwargs)
 
     # Apply the same dimension substitutions to the summand and its bindings
-    return _uxreplace.dispatch(object)(expr, {**mapper, **rule})
+    return _uxreplace_default(expr, {**mapper, **rule})
 
 
 @singledispatch
@@ -165,14 +184,15 @@ def _(iterable, rule):
 
 
 @_uxreplace_dispatch.register(dict)
+@_uxreplace_dispatch.register(frozendict)
 def _(mapper, rule):
     ret = {}
     changed = False
     for k, v in mapper.items():
         vx, flag = _uxreplace_dispatch(v, rule)
         ret[k] = vx
         changed |= flag
-    return (ret, True) if changed else (mapper, False)
+    return (mapper.__class__(ret), True) if changed else (mapper, False)
 
 
 @singledispatch
@@ -183,6 +203,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 +217,16 @@ def _(expr, args, kwargs):
     return expr.func(*args, evaluate=evaluate)
 
 
+# The SymPy bases precede UnevaluableMixin in these classes' dispatch MRO
+@_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 +288,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)
 
 
@@ -585,10 +622,14 @@ def evalrel(func=min, input=None, assumptions=None):
     return rfunc(func, *input)
 
 
-uneval_mapper = {Add: UnevalAdd, Mul: UnevalMul, Pow: UnevalPow}
+# Rounding functions also need freezing: reconstruction otherwise reevaluates
+# their lowered arguments, including during Operator unpickling.
+uneval_mapper = {Add: UnevalAdd, Mul: UnevalMul, Pow: UnevalPow,
+                 ceiling: UnevalCeil, floor: UnevalFloor}
 
 
 def unevaluate(expr):
+    """Protect lowered arithmetic from reevaluation during reconstruction."""
     if q_leaf(expr):
         return expr
 

devito/symbolics/unevaluation.py

@@ -1,7 +1,7 @@
 import sympy
 from sympy import sympify
 
-__all__ = ['Add', 'Mod', 'Mul', 'Pow']
+__all__ = ['Add', 'Mod', 'Mul', 'Pow', 'ceiling', 'floor']
 
 
 class UnevaluableMixin:
@@ -38,3 +38,11 @@ def __new__(cls, base, exp, evaluate=None, **kwargs):
             return sympy.S.One
         else:
             return cls.__base__.__new__(cls, base, exp, evaluate=False, **kwargs)
+
+
+class ceiling(sympy.ceiling, UnevaluableMixin):
+    __new__ = UnevaluableMixin.__new__
+
+
+class floor(sympy.floor, UnevaluableMixin):
+    __new__ = UnevaluableMixin.__new__