Merge pull request #2989 from devitocodes/sparse-arg-check

misc: Fix oversight regarding SparseTimeFunction index errors

Author: mloubout

devito/types/dense.py

@@ -950,16 +950,6 @@ def _arg_check(self, args, intervals, **kwargs):
         for i, s in zip(self.dimensions, data.shape, strict=True):
             i._arg_check(args, s, intervals[i])
 
-        if args.options['index-mode'] == 'int32' and \
-           args.options['linearize'] and \
-           self.is_regular 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")
-
     def _arg_finalize(self, args, alias=None):
         key = alias or self
         return {key.name: self._C_make_dataobj(alias=key, **args)}
@@ -1341,6 +1331,28 @@ def avg(self, p=None, dims=None):
         tot = self.sum(p, dims)
         return tot / len(tot.args)
 
+    def _arg_check(self, args, intervals, **kwargs):
+        """
+        Check that `args` contains legal runtime values bound to `self`.
+
+        Raises
+        ------
+        InvalidArgument
+            If an incompatibility is detected.
+        """
+        data = args[self.name]
+
+        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")
+
+        return super()._arg_check(args, intervals, **kwargs)
+
 
 class TimeFunction(Function):
 

tests/test_linearize.py

@@ -149,6 +149,18 @@ def test_interpolation_enforcing_int64_indexing():
     assert 'long p_src_stride0' in str(op)  # for `src`
 
 
+def test_interpolation_enforcing_int64_indexing_v2():
+    grid = Grid(shape=(4, 4))
+
+    # Too many elements for int32 (1e10), but should not trigger any errors, as
+    # will be indexed with int64
+    rec = SparseTimeFunction(name='rec', grid=grid, npoint=int(1e6), nt=int(1e4))
+    u = TimeFunction(name="u", grid=grid, time_order=2)
+
+    Operator(rec.interpolate(expr=u.forward),
+             opt=('advanced', {'linearize': True, 'index-mode': 'int32'}))
+
+
 def test_interpolation_msf():
     grid = Grid(shape=(4, 4))