Merge pull request #3012 from devitocodes/fix-staggered-sum-indices

api: Fix staggered sum indices

Author: mloubout

devito/core/gpu.py

@@ -162,7 +162,10 @@ def wrapper(expressions, mode='default', options=None, **kwargs1):
                 # small kernels typically generated by recursive compilation
                 par_tile0 = options0['par-tile']
                 par_tile = options.get('par-tile')
-                if par_tile0 and par_tile:
+                if par_tile is False:
+                    # The caller explicitly opted out of tiling
+                    options = {**options0, **options, 'par-tile': ParTile(None)}
+                elif par_tile0 and par_tile:
                     options = {**options0, **options, 'par-tile': par_tile}
                 elif par_tile0:
                     par_tile = ParTile(par_tile0.default, default=par_tile0.default)

devito/finite_differences/derivative.py

@@ -14,7 +14,7 @@
 from devito.warnings import warn
 
 from .differentiable import Add, Differentiable, Mul, diffify, interp_for_fd
-from .finite_difference import cross_derivative, generic_derivative
+from .finite_difference import cross_derivative, generic_derivative, indices_at
 from .rsfd import d45
 from .tools import direct, transpose
 
@@ -575,6 +575,12 @@ def _eval_fd(self, expr, **kwargs):
              shited derivative.
         - 4: Apply substitutions.
         """
+        # Differentiation is linear, and a sum of terms at different staggered
+        # locations must use it: `Add` reports its first argument's location,
+        # so `x0` would shift the other terms off the point they sat at.
+        if expr.is_Add and any(len(indices_at(expr, d)) > 1 for d in self.dims):
+            return expr.func(*[self._eval_fd(a, **kwargs) for a in expr.args])
+
         # Step 1: Evaluate non-derivative x0. We currently enforce a simple 2nd order
         # interpolation to avoid very expensive finite differences on top of it
         x0_deriv = self._filter_dims(self.x0)

devito/finite_differences/finite_difference.py

@@ -100,6 +100,29 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, side=None, **kw
     return expr
 
 
+def indices_at(expr, dim):
+    """
+    The locations `expr`'s terms sit at along `dim`.
+
+    Terms with no location of their own, a scalar say, contribute none.
+    """
+    indices = set()
+    for i in (expr.args if expr.is_Add else (expr,)):
+        try:
+            indices.add(i.indices_ref[dim])
+        except (AttributeError, KeyError, IndexError, TypeError):
+            continue
+    return indices
+
+
+def index_at(expr, dim):
+    """
+    Where `expr` sits along `dim`, or None if it does not say.
+    """
+    indices = indices_at(expr, dim)
+    return indices.pop() if len(indices) == 1 else None
+
+
 @check_input
 def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
                        coefficients='taylor', expand=True, weights=None, side=None):
@@ -139,8 +162,9 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
     if deriv_order == 1 and fd_order == 2 and side is None:
         fd_order = 1
 
-    # Zeroth order derivative is just the expression itself if not shifted
-    if deriv_order == 0 and not x0:
+    # Zeroth order is the identity when `expr` already sits at `x0`, not a
+    # stencil centred there.
+    if deriv_order == 0 and (not x0 or index_at(expr, dim) == x0.get(dim)):
         return expr
 
     # Enforce stable time coefficients

devito/ir/clusters/algorithms.py

@@ -742,11 +742,10 @@ def _normalize_reductions_dense(cluster, mapper, sregistry, platform):
             elif rhs in mapper:
                 # Seen this RHS already, so reuse the Array that was created for it
                 processed.append(e.func(lhs, mapper[rhs].indexify()))
-            elif rf and rf.is_Array and sum(flatten(rf._size_nodomain)) == 0:
-                # Special case: the RHS is an Array with no halo/padding, meaning
-                # that the written data values are contiguous in memory, hence
-                # we can simply reuse the Array itself as we're already in the
-                # desired memory layout
+            elif rf and rf.is_Array and rf._is_reduction_ready:
+                # Special case: the RHS is an Array whose written data values are
+                # contiguous in memory, hence we can simply reuse the Array
+                # itself as we're already in the desired memory layout
                 processed.append(e)
             else:
                 name = sregistry.make_name()

devito/ir/iet/visitors.py

@@ -7,6 +7,7 @@
 import ctypes
 from collections import OrderedDict
 from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
+from contextlib import suppress
 from itertools import chain, groupby
 from typing import Any, Generic, TypeVar
 
@@ -17,8 +18,8 @@
 from devito.exceptions import CompilationError
 from devito.ir.cgen.printer import get_printer
 from devito.ir.iet.nodes import (
-    BlankLine, Call, Expression, ExpressionBundle, Iteration, Lambda, ListMajor, Node,
-    Section, _same_as_before
+    BlankLine, Call, Definition, Expression, ExpressionBundle, Iteration, Lambda,
+    ListMajor, Node, Section, _same_as_before
 )
 from devito.ir.support.space import Backward
 from devito.symbolics import (
@@ -1275,6 +1276,23 @@ def visit_Call(self, o: Call, **kwargs) -> Iterator[ApplicationType]:
             except (AttributeError, TypeError):
                 yield from self._visit(i)
 
+    def visit_Definition(self, o: Definition, **kwargs) -> Iterator[ApplicationType]:
+        # The defined object carries expressions in its constructor arguments
+        # and in its initializer, both of which end up in the generated code
+        f = o.function
+        if f.is_LocalObject:
+            candidates = (*f.cargs, f.initvalue)
+        elif f.is_Array:
+            candidates = as_tuple(f.initvalue)
+        else:
+            return
+
+        for i in candidates:
+            # Not everything in there is a symbolic expression, e.g. a plain
+            # number, a string, or nothing at all
+            with suppress(AttributeError, TypeError):
+                yield from i.find(self.match)
+
 
 class IsPerfectIteration(Visitor):
 

devito/passes/iet/definitions.py

@@ -20,7 +20,7 @@
     VOID, Byref, DefFunction, FieldFromPointer, IndexedPointer, ListInitializer, SizeOf,
     as_long, pow_to_mul, unevaluate
 )
-from devito.tools import as_list, as_mapper, as_tuple, filter_sorted, flatten
+from devito.tools import as_list, as_mapper, as_tuple, filter_sorted, flatten, is_integer
 from devito.types import (
     Array, ComponentAccess, CustomDimension, DeviceMap, DeviceRM, Dimension, Eq, Symbol,
     size_t
@@ -91,6 +91,27 @@ def __init__(self, rcompile=None, sregistry=None, platform=None,
         self.sregistry = sregistry
         self.platform = platform
 
+        # Off inside the recursive compilation of a zero-init itself, which
+        # would otherwise ask for a zero-init of its own, ad infinitum
+        self.zero_init = (options or {}).get('zero-init', True)
+
+    def _zero_init(self, obj, storage):
+        """
+        The nodes zeroing `obj` upfront, if it asks for it, plus the efuncs
+        they call, if any.
+        """
+        if not (obj._is_zero_init and self.zero_init):
+            return (), ()
+
+        return self._make_zero_init(obj, storage)
+
+    def _make_zero_init(self, obj, storage):
+        """How to zero `obj`'s whole allocation, padding included."""
+        storage.include(self.langbb['header-memcpy'])
+        nbytes = SizeOf(obj._C_typedata)*as_long(obj.size)
+
+        return (self.langbb['host-memset'](obj._C_symbol, 0, nbytes),), ()
+
     def _alloc_object_on_low_lat_mem(self, site, obj, storage):
         """
         Allocate a LocalObject in the low latency memory.
@@ -172,11 +193,13 @@ def _alloc_host_array_on_high_bw_mem(self, site, obj, storage, *args):
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
         nbytes = SizeOf(obj._C_typedata)*as_long(obj.size)
-        alloc = self.langbb['host-alloc'](memptr, alignment, nbytes)
+        zeroing, efuncs = self._zero_init(obj, storage)
+        allocs = [decl, self.langbb['host-alloc'](memptr, alignment, nbytes),
+                  *zeroing]
 
         free = self.langbb['host-free'](obj._C_symbol)
 
-        storage.update(obj, site, allocs=(decl, alloc), frees=free)
+        storage.update(obj, site, allocs=tuple(allocs), frees=free, efuncs=efuncs)
 
     def _alloc_local_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
@@ -568,7 +591,7 @@ def __init__(self, options=None, **kwargs):
         self.gpu_create = options['gpu-create']
         self.gpu_place_transfers = options.get('place-transfers')
 
-        super().__init__(**kwargs)
+        super().__init__(options=options, **kwargs)
 
     def _alloc_local_array_on_high_bw_mem(self, site, obj, storage):
         """
@@ -579,11 +602,22 @@ def _alloc_local_array_on_high_bw_mem(self, site, obj, storage):
         dofree = self.langbb['device-free']
 
         nbytes = SizeOf(obj._C_typedata)*obj.size
-        init = doalloc(nbytes, deviceid, retobj=obj)
+
+        zeroing, efuncs = self._zero_init(obj, storage)
+        allocs = [doalloc(nbytes, deviceid, retobj=obj), *zeroing]
 
         free = dofree(obj._C_name, deviceid)
 
-        storage.update(obj, site, allocs=init, frees=free)
+        storage.update(obj, site, allocs=tuple(allocs), frees=free, efuncs=efuncs)
+
+    def _make_zero_init(self, obj, storage):
+        # No language here has a device-side memset, so use a kernel. It gains
+        # nothing from tiling, and nvc++ trips over the padded loop bounds
+        # when it is asked to tile them
+        efuncs, init = make_zero_init(obj, self.rcompile, self.sregistry,
+                                      options={'par-tile': False})
+
+        return (init,), efuncs
 
     def _map_array_on_high_bw_mem(self, site, obj, storage):
         """
@@ -702,18 +736,20 @@ def process(self, graph):
         self.place_casts(graph)
 
 
-def make_zero_init(obj, rcompile, sregistry):
+def make_zero_init(obj, rcompile, sregistry, options=None):
     cdims = []
-    for d, (h0, h1), s in zip(
-        obj.dimensions, obj._size_halo, obj.symbolic_shape, strict=True
+    for d, (h0, h1), (_, p1), s in zip(
+        obj.dimensions, obj._size_halo, obj._size_padding, obj.symbolic_shape,
+        strict=True
     ):
         if d.is_NonlinearDerived:
-            assert h0 == h1 == 0
+            assert h0 == h1
             m = 0
             M = s - 1
         else:
             m = d.symbolic_min - h0
-            M = d.symbolic_max + h1
+            # Object needing padding zeroing need symbolic padding
+            M = d.symbolic_max + h1 + (0 if is_integer(p1) else p1)
         cdims.append(CustomDimension(name=d.name, parent=d,
                                      symbolic_min=m, symbolic_max=M))
 
@@ -722,7 +758,8 @@ def make_zero_init(obj, rcompile, sregistry):
     else:
         eqns = [Eq(obj[cdims], 0)]
 
-    irs, byproduct = rcompile(eqns)
+    irs, byproduct = rcompile(eqns, options={'zero-init': False,
+                                             **(options or {})})
 
     init = irs.iet.body.body[0]
 

devito/passes/iet/languages/C.py

@@ -56,6 +56,8 @@ class CBB(LangBB):
             Call('free', (i,)),
         'host-free-pin': lambda i:
             Call('free', (i,)),
+        'host-memset': lambda i, j, k:
+            Call('memset', (i, j, k)),
         'alloc-global-symbol': lambda i, j, k:
             Call('memcpy', (i, j, k))
     }

devito/passes/iet/languages/CXX.py

@@ -141,6 +141,8 @@ class CXXBB(LangBB):
             Call('free', (i,)),
         'host-free-pin': lambda i:
             Call('free', (i,)),
+        'host-memset': lambda i, j, k:
+            Call('memset', (i, j, k)),
         'alloc-global-symbol': lambda i, j, k:
             Call('memcpy', (i, j, k))
     }

devito/passes/iet/misc.py

@@ -259,8 +259,9 @@ def _(expr, langbb, printer):
 
 @_lower_macro_math.register(RoundUp)
 def _(expr, langbb, printer):
-    return (('ROUND_UP(a,b)',
-             '((((a)%(b)) == 0) ? (a) : ((a) + (b) - ((a)%(b))))'),), {}
+    # Branchless: a ternary makes for a poor loop bound, and it crashes the
+    # nvc++ frontend outright when the loop is collapsed and optimized
+    return (('ROUND_UP(a,b)', '((((a) + (b) - 1)/(b))*(b))'),), {}
 
 
 @iet_pass

devito/types/basic.py

@@ -16,7 +16,7 @@
 from devito.data import default_allocator
 from devito.parameters import configuration
 from devito.tools import (
-    CustomDtype, Pickable, as_tuple, dtype_to_ctype, frozendict, memoized_meth,
+    CustomDtype, Pickable, as_tuple, dtype_to_ctype, flatten, frozendict, memoized_meth,
     sympy_mutex
 )
 from devito.types.args import ArgProvider
@@ -714,6 +714,12 @@ class AbstractFunction(sympy.Function, Basic, Pickable, Evaluable):
     effect if autopadding is disabled, which is the default behavior.
     """
 
+    _is_zero_init = False
+    """
+    Whether the entries outside `self`'s DOMAIN carry meaningful data rather
+    than scratch, in which case the whole allocation must be zeroed upfront.
+    """
+
     __rkwargs__ = ('name', 'dtype', 'grid', 'halo', 'ghost',
                    'alias', 'space', 'function', 'is_transient', 'avg_mode')
 
@@ -1351,6 +1357,14 @@ def _size_nodomain(self):
 
         return DimensionTuple(*sizes, getters=self.dimensions, left=left, right=right)
 
+    @property
+    def _is_reduction_ready(self):
+        """
+        True if a reduction over `self` may run over the whole allocated data
+        rather than over the DOMAIN alone, False otherwise.
+        """
+        return not sum(flatten(self._size_nodomain))
+
     @cached_property
     def _size_ghost(self):
         """