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))