From b69081bc4e79b3206d7e619af51d50f0f8ca6b0b Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Wed, 21 Oct 2026 16:11:45 +0100 Subject: [PATCH 1/3] compiler: Support split-buffer option at arg_check --- devito/core/cpu.py | 1 + devito/core/gpu.py | 1 + devito/types/dense.py | 18 ++++++++++++------ tests/test_linearize.py | 18 ++++++++++++++++++ 4 files changed, 32 insertions(+), 6 deletions(-) diff --git a/devito/core/cpu.py b/devito/core/cpu.py index 3435a472e6..8bd9934ad1 100644 --- a/devito/core/cpu.py +++ b/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) diff --git a/devito/core/gpu.py b/devito/core/gpu.py index 536a06422e..213b0a1cee 100644 --- a/devito/core/gpu.py +++ b/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) diff --git a/devito/types/dense.py b/devito/types/dense.py index 30f543a47c..96c99eb33b 100644 --- a/devito/types/dense.py +++ b/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) diff --git a/tests/test_linearize.py b/tests/test_linearize.py index c78880924d..272d8e1054 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -6,6 +6,7 @@ Eq, Function, Grid, Inc, MatrixSparseTimeFunction, Operator, SparseTimeFunction, TimeFunction, configuration, sin, switchconfig ) +from devito.exceptions import InvalidArgument from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry from devito.passes import Graph, generate_macros, linearize from devito.types import Array, Bundle, CustomDimension, DefaultDimension @@ -161,6 +162,23 @@ def test_interpolation_enforcing_int64_indexing_v2(): opt=('advanced', {'linearize': True, 'index-mode': 'int32'})) +@switchconfig(autopadding=False) +@pytest.mark.parametrize('index_mode', ['int32', 'int64']) +def test_large_allocation(index_mode): + grid = Grid(shape=(1024, 1024, 1024)) + u = TimeFunction(name='u', grid=grid, time_order=2, space_order=0) + op = Operator(Eq(u.forward, u + 1), name='large', + opt=('advanced', {'linearize': True, 'index-mode': index_mode})) + + # Validate the allocation size without allocating or accessing its storage + data = np.broadcast_to(np.zeros(1, dtype=u.dtype), u.shape_allocated) + if index_mode == 'int32': + with pytest.raises(InvalidArgument, match='int32 pointer arithmetic'): + op.arguments(u=data, time_m=0, time_M=2) + else: + op.arguments(u=data, time_m=0, time_M=2) + + def test_interpolation_msf(): grid = Grid(shape=(4, 4)) From e57fb35e8d3439be5121cd27ac8dd427591737de Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Thu, 22 Oct 2026 14:33:43 +0100 Subject: [PATCH 2/3] compiler: Bind new callable parameters at each call site --- devito/ir/iet/nodes.py | 8 ++++++-- devito/ir/iet/visitors.py | 4 ++-- devito/passes/iet/engine.py | 27 ++++++++++++++++++------- devito/symbolics/extended_sympy.py | 2 ++ devito/types/array.py | 8 +++++--- tests/test_iet.py | 32 ++++++++++++++++++++++++++---- tests/test_symbolics.py | 12 ++++++++--- 7 files changed, 72 insertions(+), 21 deletions(-) diff --git a/devito/ir/iet/nodes.py b/devito/ir/iet/nodes.py index 7c12b39647..f340b0fea5 100644 --- a/devito/ir/iet/nodes.py +++ b/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) diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 6186ff4e76..7c68ee75f2 100644 --- a/devito/ir/iet/visitors.py +++ b/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 ( @@ -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) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index fe9a78b65b..9ddc8714fa 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -14,15 +14,16 @@ 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 +135,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 +171,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): @@ -723,7 +728,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 +780,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 +814,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: diff --git a/devito/symbolics/extended_sympy.py b/devito/symbolics/extended_sympy.py index eeb290feee..4f8df6d914 100644 --- a/devito/symbolics/extended_sympy.py +++ b/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}" diff --git a/devito/types/array.py b/devito/types/array.py index 1988e2b0a8..98f8932813 100644 --- a/devito/types/array.py +++ b/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") diff --git a/tests/test_iet.py b/tests/test_iet.py index 99a94ed24c..41caae0f84 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -19,8 +19,8 @@ from devito.passes.iet.engine import Graph, reuse_efuncs from devito.passes.iet.languages.C import CDataManager 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 @@ -430,17 +430,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(): @@ -542,6 +545,27 @@ def test_codegen_quality0(): assert foo1.parameters[0] is a +def test_graph_parameter_bindings(): + x, y, a, b, p = [Symbol(name=n) for n in ('x', 'y', 'a', 'b', 'p')] + inner = Callable('inner', Call('external', (x, y)), 'void', parameters=(x, y)) + root = make_callable('root', [Call('inner', (a + 1, b)), + Call('inner', (b, a + 2))]) + graph = Graph(root) + graph.apply(lambda iet: (iet, {'efuncs': [inner]})) + + def transform(iet): + if iet.name != 'inner': + return iet, {} + return iet._rebuild(body=Call('external', (p,))), {'args': {p: x + y}} + + graph.apply(transform) + + assert graph.efuncs['inner'].parameters == (p,) + assert [c.arguments for c in FindNodes(Call).visit(graph.root)] == [ + (a + b + 1,), (a + b + 2,) + ] + + def test_reuse_efuncs_natural_numbering(): def make_foo(name, start): diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 51a759a705..431b96c424 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -16,7 +16,8 @@ from devito.mpi.halo_scheme import HaloTouch from devito.symbolics import ( # noqa INT, BaseCast, CallFromPointer, Cast, DefFunction, FieldFromComposite, - FieldFromPointer, IntDiv, ListInitializer, Namespace, ReservedWord, RoundUp, Rvalue, + FieldFromPointer, IndexedPointer, IntDiv, ListInitializer, Namespace, + ReservedWord, RoundUp, Rvalue, SizeOf, VectorAccess, evalrel, pow_to_mul, q_negative, q_positive, retrieve_derivatives, retrieve_functions, retrieve_indexed, subs_if_composite, uxreplace, xreplace_indices @@ -645,11 +646,16 @@ def test_component_access(): assert cf2 == cf1 -def test_component_access_symbol_printing(): +@pytest.mark.parametrize('pointer,expected', [ + (False, '1.250F*acc.x'), (True, '1.250F*acc[0].x') +]) +def test_component_access_symbol_printing(pointer, expected): acc = dSymbol(name='acc', dtype=dtypes_vector_mapper[(np.float32, 4)]) + if pointer: + acc = IndexedPointer(acc, 0) expr = ComponentAccess(acc, 0) - assert ccode(1.25*expr, dtype=expr.dtype) == '1.250F*acc.x' + assert ccode(1.25*expr, dtype=expr.dtype) == expected def test_vector_access(): From 3d4dbd63e730672aafb737a3f778d3191dc1daa1 Mon Sep 17 00:00:00 2001 From: Fabio Luporini Date: Thu, 22 Oct 2026 16:28:01 +0100 Subject: [PATCH 3/3] compiler: Preserve types and options for split buffers --- devito/core/gpu.py | 2 ++ devito/ir/iet/visitors.py | 2 +- devito/passes/iet/engine.py | 5 +++-- devito/types/basic.py | 13 ++++++------- tests/test_iet.py | 10 ++++++++++ tests/test_linearize.py | 12 +++++++----- tests/test_symbolics.py | 5 ++--- 7 files changed, 31 insertions(+), 18 deletions(-) diff --git a/devito/core/gpu.py b/devito/core/gpu.py index 213b0a1cee..7fe399771f 100644 --- a/devito/core/gpu.py +++ b/devito/core/gpu.py @@ -152,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', diff --git a/devito/ir/iet/visitors.py b/devito/ir/iet/visitors.py index 7c68ee75f2..4b83d713a5 100644 --- a/devito/ir/iet/visitors.py +++ b/devito/ir/iet/visitors.py @@ -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 diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index 9ddc8714fa..3d10296824 100644 --- a/devito/passes/iet/engine.py +++ b/devito/passes/iet/engine.py @@ -14,8 +14,9 @@ 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, uxreplace) +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 diff --git a/devito/types/basic.py b/devito/types/basic.py index fc8940f5e9..bb3b31a922 100644 --- a/devito/types/basic.py +++ b/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 diff --git a/tests/test_iet.py b/tests/test_iet.py index 41caae0f84..90183abe7d 100644 --- a/tests/test_iet.py +++ b/tests/test_iet.py @@ -545,6 +545,16 @@ 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) + + def test_graph_parameter_bindings(): x, y, a, b, p = [Symbol(name=n) for n in ('x', 'y', 'a', 'b', 'p')] inner = Callable('inner', Call('external', (x, y)), 'void', parameters=(x, y)) diff --git a/tests/test_linearize.py b/tests/test_linearize.py index 272d8e1054..cfd019f9c8 100644 --- a/tests/test_linearize.py +++ b/tests/test_linearize.py @@ -6,6 +6,7 @@ Eq, Function, Grid, Inc, MatrixSparseTimeFunction, Operator, SparseTimeFunction, TimeFunction, configuration, sin, switchconfig ) +from devito.data.allocators import DataReference from devito.exceptions import InvalidArgument from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry from devito.passes import Graph, generate_macros, linearize @@ -166,17 +167,18 @@ def test_interpolation_enforcing_int64_indexing_v2(): @pytest.mark.parametrize('index_mode', ['int32', 'int64']) def test_large_allocation(index_mode): grid = Grid(shape=(1024, 1024, 1024)) - u = TimeFunction(name='u', grid=grid, time_order=2, space_order=0) + # Validate the allocation size without allocating or accessing its storage + data = np.broadcast_to(np.zeros(1, dtype=np.float32), (3, *grid.shape)) + u = TimeFunction(name='u', grid=grid, time_order=2, space_order=0, + allocator=DataReference(data)) op = Operator(Eq(u.forward, u + 1), name='large', opt=('advanced', {'linearize': True, 'index-mode': index_mode})) - # Validate the allocation size without allocating or accessing its storage - data = np.broadcast_to(np.zeros(1, dtype=u.dtype), u.shape_allocated) if index_mode == 'int32': with pytest.raises(InvalidArgument, match='int32 pointer arithmetic'): - op.arguments(u=data, time_m=0, time_M=2) + op.arguments(time_m=0, time_M=2) else: - op.arguments(u=data, time_m=0, time_M=2) + op.arguments(time_m=0, time_M=2) def test_interpolation_msf(): diff --git a/tests/test_symbolics.py b/tests/test_symbolics.py index 431b96c424..f68da12f20 100644 --- a/tests/test_symbolics.py +++ b/tests/test_symbolics.py @@ -16,9 +16,8 @@ from devito.mpi.halo_scheme import HaloTouch from devito.symbolics import ( # noqa INT, BaseCast, CallFromPointer, Cast, DefFunction, FieldFromComposite, - FieldFromPointer, IndexedPointer, IntDiv, ListInitializer, Namespace, - ReservedWord, RoundUp, Rvalue, - SizeOf, VectorAccess, evalrel, pow_to_mul, q_negative, q_positive, + FieldFromPointer, IndexedPointer, IntDiv, ListInitializer, Namespace, ReservedWord, + RoundUp, Rvalue, SizeOf, VectorAccess, evalrel, pow_to_mul, q_negative, q_positive, retrieve_derivatives, retrieve_functions, retrieve_indexed, subs_if_composite, uxreplace, xreplace_indices )