Merge pull request #3042 from devitocodes/temp_alloc_cast

compiler: Cast operands to long when calculating sizes for array allocation

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -601,7 +601,7 @@ def _alloc_local_array_on_high_bw_mem(self, site, obj, storage):
         doalloc = self.langbb['device-alloc']
         dofree = self.langbb['device-free']
 
-        nbytes = SizeOf(obj._C_typedata)*obj.size
+        nbytes = SizeOf(obj._C_typedata)*as_long(obj.size)
 
         zeroing, efuncs = self._zero_init(obj, storage)
         allocs = [doalloc(nbytes, deviceid, retobj=obj), *zeroing]

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:
 

tests/test_gpu_openmp.py

@@ -217,7 +217,8 @@ def test_array_rw(self):
         assert len(op.body.allocs) == 2
         assert str(op.body.allocs[0]) ==\
             ('float * r0_vec = (float *)'
-             'omp_target_alloc(x_size*y_size*z_size*sizeof(float),'
+             'omp_target_alloc('
+             'sizeof(float)*(long)z_size*(long)y_size*(long)x_size,'
              'omp_get_default_device());')
         assert str(op.body.allocs[1]) ==\
             'init0(x_M,x_m,y_M,y_m,z_M,z_m,r0_vec,y_size,z_size);'