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/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
 )
@@ -394,7 +394,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)

devito/passes/iet/engine.py

@@ -21,8 +21,8 @@
 )
 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
@@ -781,7 +781,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
@@ -790,29 +791,21 @@ def _update_args_efunc(root, efuncs, dag):
     if not (new_params or drop_params):
         return efuncs
 
-    # Create the new parameters and arguments lists
-
-    def _filter(v, efunc=None):
-        processed = [a for i, a in enumerate(v) if i not in drop_params]
+    # Decide additions from the signature, since an existing Call argument may
+    # bind the same symbol to a different parameter.
+    new_params = [a for a in new_params if a not in root.parameters]
 
-        for a in new_params:
-            if a in processed:
-                # A child efunc trying to add a symbol already added by a
-                # sibling efunc
-                continue
+    if isinstance(root, EntryFunction):
+        # Temporaries (e.g., Arrays) *cannot* be arguments of an EntryFunction.
+        # These remain for now undefined inside `root`.
+        new_params = [a for a in new_params if isinstance(a, ArgProvider)]
 
-            if isinstance(efunc, EntryFunction) and not isinstance(a, ArgProvider):
-                # Temporaries (e.g., Arrays) *cannot* be arguments of an
-                # EntryFunction. So if we end up here, `a` remains for
-                # now undefined inside `efunc`
-                continue
-
-            processed.append(a)
-
-        return processed
+    # Create the new parameters and arguments lists
+    def _filter(v):
+        return [a for i, a in enumerate(v) if i not in drop_params] + new_params
 
     efuncs = dict(efuncs)
-    efuncs[root.name] = root._rebuild(parameters=_filter(root.parameters, root))
+    efuncs[root.name] = root._rebuild(parameters=_filter(root.parameters))
 
     # Update all call sites to use the new signature
     for n in dag.downstream(root.name):

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
@@ -372,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}"

devito/types/array.py

@@ -645,9 +645,11 @@ class ComponentAccess(Expr, Pickable):
     __rkwargs__ = ('index',)
 
     def __new__(cls, arg, index=0, **kwargs):
-        if not (arg.is_Symbol or arg.is_Indexed):
-            raise ValueError(f"Expected Symbol or Indexed, got `{type(arg)}` "
-                             "instead")
+        from devito.symbolics import IndexedPointer
+
+        if not (arg.is_Symbol or arg.is_Indexed or isinstance(arg, IndexedPointer)):
+            raise ValueError(f"Expected Symbol, Indexed or IndexedPointer, "
+                             f"got `{type(arg)}` instead")
         if not is_integer(index) or index > 3:
             raise ValueError("Expected 0 <= index < 4")
 

devito/types/basic.py

@@ -58,6 +58,12 @@ class CodeSymbol:
     for a given architecture.
     """
 
+    _C_modifier = None
+    """
+    A suffix added to the type in a function signature. E.g., `&` for a
+    reference parameter, or ` const` for a constant pointer.
+    """
+
     @abc.abstractmethod
     def __init__(self, *args, **kwargs):
         return
@@ -2057,13 +2063,6 @@ def _mem_internal_eager(self):
     def _mem_internal_lazy(self):
         return self._liveness == 'lazy'
 
-    """
-    A modifier added to the declaration of the LocalType when it appears in a
-    function signature. For example, a subclass might define `_C_modifier = '&'`
-    to impose pass-by-reference semantics.
-    """
-    _C_modifier = None
-
     """
     One or more optional keywords added to the declaration of the LocalType
     in between the type and the variable name when it appears in a function

devito/types/dense.py

@@ -23,7 +23,7 @@
 from devito.symbolics import FieldFromPointer, IndexedPointer, normalize_args
 from devito.tools import (
     ReducerMap, as_tuple, c_restrict_void_p, dtype_to_ctype, flatten, humanbytes,
-    is_integer, memoized_meth, mpi4py_mapper
+    is_integer, memoized_meth, mpi4py_mapper, prod
 )
 from devito.types.args import ArgProvider
 from devito.types.basic import AbstractFunction
@@ -1351,11 +1351,17 @@ def _arg_check(self, args, intervals, **kwargs):
         if args.options['index-mode'] == 'int32' and \
            args.options['linearize'] and \
            data.size - 1 >= np.iinfo(np.int32).max:
-            raise InvalidArgument(f"`{self.name}`, with its {data.size} elements, is too "
-                                  "big for int32 pointer arithmetic. Consider using the "
-                                  "'index-mode=int64' option, the save=Buffer(..) "
-                                  "API (TimeFunction only), or domain "
-                                  "decomposition via MPI")
+            size = data.size
+            if args.options['split-modulo-buffers'] and len(data.shape) == self.ndim:
+                size = prod(s for d, s in zip(self.dimensions, data.shape, strict=True)
+                            if not d.is_Stepping)
+            if size - 1 >= np.iinfo(np.int32).max:
+                raise InvalidArgument(
+                    f"`{self.name}`, with its {size} indexed elements, is too "
+                    "big for int32 pointer arithmetic. Consider using the "
+                    "'index-mode=int64' option, the save=Buffer(..) "
+                    "API (TimeFunction only), or domain decomposition via MPI"
+                )
 
         return super()._arg_check(args, intervals, **kwargs)
 

devito/types/misc.py

@@ -110,6 +110,11 @@ def __new__(cls, base, *args, strides_map=None, accessor=None):
 
         return obj
 
+    @classmethod
+    def class_key(cls):
+        # Keep arithmetic ordering consistent with unlinearized accesses
+        return Indexed.class_key()
+
     def __repr__(self):
         return "{}({})".format(self.name, ", ".join(str(i) for i in self.indices))
 

tests/test_iet.py

@@ -6,21 +6,22 @@
 import sympy
 
 from devito import (  # noqa
-    Dimension, Eq, Function, Grid, Operator, TimeFunction, switchconfig
+    Dimension, Eq, Function, Grid, Operator, SubDimension, TimeFunction, switchconfig
 )
 from devito.ir import SymbolRegistry
 from devito.ir.iet import (
-    Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall, DummyExpr,
-    ElementalFunction, FindNodes, FindSymbols, Iteration, KernelLaunch, Lambda, List,
-    Switch, Transformer, Uxreplace, filter_iterations, make_callable, make_efunc,
-    retrieve_iteration_tree
+    Call, Callable, CGen, Conditional, Definition, Dereference, DeviceCall,
+    DeviceFunction, DummyExpr, ElementalFunction, EntryFunction, FindNodes, FindSymbols,
+    Iteration, KernelLaunch, Lambda, List, Switch, Transformer, Uxreplace,
+    filter_iterations, make_callable, make_efunc, retrieve_iteration_tree
 )
 from devito.ir.iet.visitors import sorted_efuncs
 from devito.passes.iet.engine import Graph, reuse_efuncs
 from devito.passes.iet.languages.C import CDataManager
+from devito.passes.iet.misc import abridge_dim_names
 from devito.symbolics import (
-    FLOAT, Byref, Class, FieldFromComposite, InlineIf, ListInitializer, Macro, SizeOf,
-    String
+    FLOAT, Byref, Class, FieldFromComposite, IndexedPointer, InlineIf, ListInitializer,
+    Macro, SizeOf, String
 )
 from devito.symbolics.extended_dtypes import c_complex
 from devito.tools import CustomDtype, as_tuple, dtype_to_ctype
@@ -444,17 +445,20 @@ def test_call_indexed():
 }"""
 
 
-def test_call_retobj_indexed():
+@pytest.mark.parametrize('pointer', [False, True])
+def test_call_retobj_indexed(pointer):
     grid = Grid(shape=(10, 10))
 
     u = Function(name='u', grid=grid)
     v = Function(name='v', grid=grid)
 
-    call = Call('foo', [u], retobj=v.indexify())
+    retobj = IndexedPointer(v.indexed, grid.dimensions) if pointer else v.indexify()
+    call = Call('foo', [u], retobj=retobj)
 
     assert str(call) == "v[x][y] = foo(u_vec);"
 
     assert not call.defines
+    assert set(grid.dimensions).issubset(call.expr_symbols)
 
 
 def test_call_lambda_transform():
@@ -556,6 +560,78 @@ def test_codegen_quality0():
     assert foo1.parameters[0] is a
 
 
+def test_const_pointer_parameter():
+    class ConstPointer(Symbol):
+        _C_modifier = ' const'
+
+    p = ConstPointer(name='p', dtype=POINTER(dtype_to_ctype(np.float32)))
+    foo = Callable('foo', Call('external', (p,)), 'void', parameters=(p,))
+
+    assert 'void foo(float * const p)' in str(foo)
+
+
+@pytest.mark.parametrize('local', [False, True])
+@pytest.mark.parametrize('collision', [False, True])
+def test_graph_parameters(local, collision):
+    x, y, a, b, c, d, p = [Symbol(name=n) for n in ('x', 'y', 'a', 'b',
+                                                 'c', 'd', 'p')]
+    inner = Callable('inner', Call('external', (x, y)), 'void', parameters=(x, y))
+    calls = [Call('inner', (a + 1, b)),
+             Call('inner', (p if collision else b, a + 2))]
+    body = [DummyExpr(p, a + b, init=True), *calls] if local else calls
+    middle = Callable('middle', body, 'void', parameters=(a, b))
+    root = Callable('root', [Call('middle', (c + 1, d)),
+                             Call('middle', (d, c + 2))],
+                    'void', parameters=(c, d))
+    graph = Graph(root)
+    graph.efuncs.update({i.name: i for i in (middle, inner)})
+
+    def transform(iet):
+        if iet.name != 'inner':
+            return iet, {}
+        return iet._rebuild(body=Call('external', (x, p))), {}
+
+    graph.apply(transform)
+
+    assert graph.efuncs['inner'].parameters == (x, p)
+    middle = graph.efuncs['middle']
+    if local:
+        assert middle.parameters == (a, b)
+        expected = [(c + 1, d), (d, c + 2)]
+    elif collision:
+        assert middle.parameters == (a, p)
+        expected = [(c + 1, p), (d, p)]
+    else:
+        assert middle.parameters == (a, b, p)
+        expected = [(c + 1, d, p), (d, c + 2, p)]
+    assert [i.arguments for i in FindNodes(Call).visit(middle)] == [
+        (a + 1, p), (p if collision else b, p)
+    ]
+    assert graph.root.parameters == ((c, d) if local else (c, d, p))
+    assert [i.arguments for i in FindNodes(Call).visit(graph.root)] == expected
+
+
+@pytest.mark.parametrize('ftype', [EntryFunction, DeviceFunction])
+@pytest.mark.parametrize('pointer', [False, True])
+@pytest.mark.parametrize('clash', [False, True])
+def test_abridge_pointer_dimensions(ftype, pointer, clash):
+    x = Dimension(name='x')
+    xi = SubDimension.left(name='xi', parent=x, thickness=4)
+    a = Array(name='a', dimensions=(x,), shape=(16,))
+    index = xi + x if clash else xi
+    access = IndexedPointer(a, index) if pointer else a[index]
+    loop = Iteration(DummyExpr(access, access + 1), xi,
+                     (xi.symbolic_min, xi.symbolic_max, 1))
+    body = Iteration(loop, x, (0, 15, 1)) if clash else loop
+    foo = ftype('foo', body, 'void', parameters=(a,))
+
+    foo = abridge_dim_names(foo)
+
+    inner = FindNodes(Iteration).visit(foo)[-1]
+    assert inner.dim.name == ('xi' if clash else 'x')
+    assert inner.limits == loop.limits
+
+
 @pytest.mark.parametrize('copy_symbols', [False, True])
 def test_reuse_efuncs_natural_numbering(copy_symbols):