compiler: Support split-buffer option at arg_check

devito/core/cpu.py

@@ -97,6 +97,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['linearize'] = oo.pop('linearize', cls.LINEARIZE)
         o['mapify-reduce'] = oo.pop('mapify-reduce', cls.MAPIFY_REDUCE)
         o['index-mode'] = oo.pop('index-mode', cls.INDEX_MODE)
+        o['split-modulo-buffers'] = False
         o['place-transfers'] = oo.pop('place-transfers', True)
         o['errctl'] = oo.pop('errctl', cls.ERRCTL)
         o['scalar-min-type'] = oo.pop('scalar-min-type', cls.SCALAR_MIN_TYPE)

devito/core/gpu.py

@@ -115,6 +115,7 @@ def _normalize_kwargs(cls, **kwargs):
         o['linearize'] = oo.pop('linearize', False)
         o['mapify-reduce'] = oo.pop('mapify-reduce', cls.MAPIFY_REDUCE)
         o['index-mode'] = oo.pop('index-mode', cls.INDEX_MODE)
+        o['split-modulo-buffers'] = False
         o['place-transfers'] = oo.pop('place-transfers', True)
         o['errctl'] = oo.pop('errctl', cls.ERRCTL)
         o['scalar-min-type'] = oo.pop('scalar-min-type', cls.SCALAR_MIN_TYPE)
@@ -151,6 +152,8 @@ def wrapper(expressions, mode='default', options=None, **kwargs1):
             options = options or {}
 
             if mode == 'host':
+                options = {'split-modulo-buffers': options0['split-modulo-buffers'],
+                           **options}
                 target = {
                     'platform': 'cpu64',
                     'language': 'C' if options0['par-disabled'] else 'openmp',

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/nodes.py

@@ -18,7 +18,7 @@
     AFFINE, INBOUND, PARALLEL, PARALLEL_IF_ATOMIC, PARALLEL_IF_PVT, SEQUENTIAL,
     VECTORIZED, Forward, PrefetchUpdate, Property, WithLock
 )
-from devito.symbolics import CallFromPointer, ListInitializer
+from devito.symbolics import CallFromPointer, IndexedPointer, ListInitializer
 from devito.tools import (
     Signer, as_tuple, ctypes_to_cstr, filter_ordered, filter_sorted, flatten
 )
@@ -344,6 +344,10 @@ def _rebuild(self, *args, **kwargs):
             kwargs['arguments'] = [mapper.get(i, i) for i in self.arguments]
         return super()._rebuild(**kwargs)
 
+    def rebuild_args(self, func):
+        """Rebuild the Call, applying `func` to its callee's arguments."""
+        return self._rebuild(arguments=func(self.arguments))
+
     @property
     def children(self):
         return tuple(i for i in self.arguments if isinstance(i, (Call, Lambda)))
@@ -394,7 +398,7 @@ def expr_symbols(self):
         if self.base is not None:
             retval.append(self.base)
 
-        if isinstance(self.retobj, Indexed):
+        if isinstance(self.retobj, (Indexed, IndexedPointer)):
             retval.extend(self.retobj.free_symbols)
         elif self.retobj is not None:
             retval.append(self.retobj)

devito/ir/iet/visitors.py

@@ -23,7 +23,7 @@
 )
 from devito.ir.support.space import Backward
 from devito.symbolics import (
-    FieldFromComposite, FieldFromPointer, ListInitializer, uxreplace
+    FieldFromComposite, FieldFromPointer, IndexedPointer, ListInitializer, uxreplace
 )
 from devito.symbolics.extended_dtypes import NoDeclStruct
 from devito.tools import (
@@ -342,7 +342,7 @@ def _gen_value(self, obj, mode=1, masked=()):
                 strtype = f'{strtype}{self._restrict_keyword}'
         strtype = ' '.join(qualifiers + [strtype])
 
-        if obj.is_LocalObject and obj._C_modifier is not None and mode == 2:
+        if obj._C_modifier is not None and mode == 2:
             strtype += obj._C_modifier
 
         strname = obj._C_name
@@ -615,7 +615,7 @@ def visit_Call(self, o, nested_call=False):
             call = MultilineCall(o.name, arguments, True, o.is_indirect, cast,
                                  o.templates)
             if retobj.is_Indexed or \
-               isinstance(retobj, (FieldFromComposite, FieldFromPointer)):
+               isinstance(retobj, (FieldFromComposite, FieldFromPointer, IndexedPointer)):
                 return c.Assign(self.ccode(retobj), call)
             else:
                 return c.Initializer(c.Value(rettype, retobj._C_name), call)
@@ -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
@@ -477,6 +479,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):
@@ -520,15 +527,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

@@ -14,15 +14,17 @@
 from devito.mpi.distributed import MPINeighborhood
 from devito.mpi.routines import Gather, HaloUpdate, HaloWait, MPIMsg, Scatter
 from devito.passes import needs_transfer
-from devito.symbolics import FieldFromComposite, FieldFromPointer, IndexedPointer, search
+from devito.symbolics import (
+    FieldFromComposite, FieldFromPointer, IndexedPointer, search, uxreplace
+)
 from devito.tools import (
     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,
-    Indirection, ModuloDimension, NPThreads, NThreadsBase, Pointer, SharedData, Symbol,
-    Temp, ThreadArray, Wildcard
+    IndexedData, Indirection, ModuloDimension, NPThreads, NThreadsBase, Pointer,
+    SharedData, Symbol, Temp, ThreadArray, Wildcard
 )
 from devito.types.args import ArgProvider
 from devito.types.dense import DiscreteFunction
@@ -134,6 +136,10 @@ def apply(self, func, **kwargs):
 
         Callable parameters and Call arguments are reconciled before the graph
         walk, after each changed node, and after the pass has completed.
+
+        The returned metadata may include `args`, a mapping from new parameters
+        to expressions in the original parameters. Each call site instantiates
+        these expressions using its own argument bindings.
         """
         _update_args(self)
 
@@ -166,7 +172,7 @@ def apply(self, func, **kwargs):
             efuncs[i] = efunc
             efuncs.update(dict([(i.name, i) for i in new_efuncs]))
 
-            efuncs = _update_args_efunc(efunc, efuncs, dag)
+            efuncs = _update_args_efunc(efunc, efuncs, dag, metadata.get('args'))
 
         # Minimize code size
         if len(efuncs) > len(self.efuncs):
@@ -564,6 +570,8 @@ def abstract_objects(objects0):
     mapper = {}
     sregistry = SymbolRegistry()
     for i in objects:
+        if i in mapper:
+            continue
         abstract_object(i, mapper, sregistry)
 
     return mapper
@@ -723,7 +731,7 @@ def _(i, mapper, sregistry):
     mapper[i] = i._rebuild(name=sregistry.make_name(prefix='nthreads'))
 
 
-def _update_args_efunc(root, efuncs, dag):
+def _update_args_efunc(root, efuncs, dag, args=None):
     """
     Re-derive the parameters of `root` and apply the changes in cascade through
     the `efuncs`.
@@ -775,7 +783,8 @@ def _update_args_efunc(root, efuncs, dag):
     # 4) removed a function that was previously necessary
     functions = FindSymbols('symbolics').visit(root.body)
     drop_params.extend(a for a in root.parameters
-                       if a.is_AbstractFunction and a not in functions)
+                       if (a.is_AbstractFunction or isinstance(a, IndexedData))
+                       and a.function not in functions)
 
     # Must record the index, not the param itself, since a param may be
     # bound to whatever arg, possibly a generic SymPy expr
@@ -808,9 +817,16 @@ def _filter(v, efunc=None):
     efuncs = dict(efuncs)
     efuncs[root.name] = root._rebuild(parameters=_filter(root.parameters, root))
 
+    def _filter_call(values):
+        mapper = {}
+        if args:
+            bindings = dict(zip(root.parameters, values, strict=True))
+            mapper = {p: uxreplace(expr, bindings) for p, expr in args.items()}
+        return [mapper.get(a, a) for a in _filter(values)]
+
     # Update all call sites to use the new signature
     for n in dag.downstream(root.name):
-        mapper = {c: c._rebuild(arguments=_filter(c.arguments))
+        mapper = {c: c.rebuild_args(_filter_call)
                   for c in FindNodes(Call).visit(efuncs[n])
                   if c.name == root.name}
         if not mapper:

devito/passes/iet/misc.py

@@ -8,8 +8,7 @@
 from devito.finite_differences.differentiable import SafeInv
 from devito.ir import (
     Any, DummyExpr, EmptyList, FindApplications, FindNodes, FindSymbols, Forward,
-    Iteration, Prodder, Transformer, Uxreplace, filter_iterations, pull_dims,
-    retrieve_iteration_tree
+    Iteration, Prodder, Transformer, Uxreplace, filter_iterations, retrieve_iteration_tree
 )
 from devito.ir.iet.efunc import DeviceFunction, EntryFunction
 from devito.passes.iet.engine import iet_pass
@@ -188,8 +187,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 +199,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)
 
@@ -371,10 +371,8 @@ def _(iet):
 
 
 def _rename_subdims(target, dimensions):
-    # Find SubDimensions or SubDimension-derived dimensions used as indices in
-    # the expression
-    indexeds = FindSymbols('indexeds').visit(target)
-    dims = pull_dims(indexeds, flag=False)
+    # Find SubDimensions or SubDimension-derived dimensions used in expressions
+    dims = FindSymbols('dimensions').visit(target)
     dims = [d for d in dims if any(dim.is_AbstractSub for dim in d._defines)]
     dims = [d for d in dims if not d.is_SubIterator]
     names = [d.root.name for d in dims]

devito/symbolics/extended_sympy.py

@@ -547,6 +547,8 @@ def base(self):
     def index(self):
         return self._index
 
+    indices = index
+
     def __str__(self):
         indices = ''.join(f'[{i}]' for i in self.index)
         return f"{self.base}{indices}"