tests: Add test for temp array allocation

Author: EdCaunt

devito/passes/iet/langbase.py

@@ -538,6 +538,7 @@ def __new__(cls, function, *args):
 
     @property
     def size(self):
+        # TODO: Do operands need casting to int64 here?
         return prod(s for _, s in self)
 
     @property

devito/types/basic.py

@@ -1141,13 +1141,6 @@ def ndim(self):
         """The rank of the object."""
         return len(self.indices)
 
-    # @property
-    # def safe_size(self):
-    #     """
-    #     The symbolic size of the object, casting to LONG to prevent overflow.
-    #     """
-    #     return prod(as_long(i) for i in self.symbolic_shape)
-
     @property
     def symbolic_shape(self):
         """

tests/test_dse.py

@@ -2976,6 +2976,30 @@ def test_subdims_wo_block_temps(self):
 
         op.apply()
 
+    @pytest.mark.parametrize('lang', ['C', 'CXX'])
+    def test_temp_array_allocation_uses_long(self, lang):
+        """
+        Ensure that operands are cast to long when calculating the size of a
+        temporary array to prevent int32 wraparound.
+        """
+        grid = Grid(shape=(31, 17, 25))
+
+        a = Function(name='a', grid=grid)
+        f = TimeFunction(name='f', grid=grid)
+        g = TimeFunction(name='g', grid=grid)
+
+        eq0 = Eq(f.forward, g + sin(a).dx)
+        eq1 = Eq(g.forward, f + sin(a).dx)
+
+        opt = ('advanced', {'cire-minmem': True, 'index-mode': 'int64'})
+        with switchconfig(language=lang):
+            op = Operator([eq0, eq1], opt=opt)
+
+            cast = '(long)' if lang == 'C' else 'static_cast<long>'
+
+            # 3 casts in memalign, 3 in memset
+            assert str(op).count(cast) == 6
+
 
 class TestIsoAcoustic: