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)

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)
 

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