compiler: fix temp allocation with long index-mode

Author: mloubout

devito/ir/cgen/printer.py

@@ -370,7 +370,7 @@ def _print_UnaryOp(self, expr, op=None, parenthesize=False):
 
     def _print_Cast(self, expr):
         cast = f'({self._print(expr._C_ctype)}{self._print(expr.stars)})'
-        return self._print_UnaryOp(expr, op=cast)
+        return self._print_UnaryOp(expr, op=cast, parenthesize=True)
 
     def _print_ComponentAccess(self, expr):
         return f"{self._print(expr.base)}.{expr.sindex}"

devito/passes/iet/definitions.py

@@ -19,7 +19,7 @@
 from devito.passes.iet.langbase import LangBB
 from devito.symbolics import (
     Byref, DefFunction, FieldFromPointer, IndexedPointer, ListInitializer,
-    SizeOf, VOID, pow_to_mul, unevaluate
+    SizeOf, VOID, pow_to_mul, unevaluate, LONG, retrieve_symbols
 )
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import (
@@ -90,6 +90,17 @@ def __init__(self, rcompile=None, sregistry=None, platform=None, **kwargs):
         self.rcompile = rcompile
         self.sregistry = sregistry
         self.platform = platform
+        self.index_mode = kwargs.get('options', {'index-mode': 'int32'})['index-mode']
+
+    def intm(self, nbytes):
+        if self.index_mode == 'int64':
+            try:
+                syms = retrieve_symbols(nbytes)
+                return nbytes.subs({s: LONG(s) for s in syms})
+            except AttributeError:
+                return LONG(nbytes)
+        else:
+            return nbytes
 
     def _alloc_object_on_low_lat_mem(self, site, obj, storage):
         """
@@ -136,7 +147,7 @@ def _alloc_array_on_global_mem(self, site, obj, storage):
 
         # Copy input array into global array
         name = self.sregistry.make_name(prefix='init_global')
-        nbytes = SizeOf(obj._C_typedata)*obj.size
+        nbytes = SizeOf(obj._C_typedata)*self.intm(obj.size)
         body = [Definition(src),
                 self.langbb['alloc-global-symbol'](obj.indexed, src.indexed, nbytes)]
         efunc = make_callable(name, body)
@@ -159,7 +170,7 @@ def _alloc_host_array_on_high_bw_mem(self, site, obj, storage, *args):
 
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
-        nbytes = SizeOf(obj._C_typedata)*obj.size
+        nbytes = SizeOf(obj._C_typedata)*self.intm(obj.size)
         alloc = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         free = self.langbb['host-free'](obj._C_symbol)
@@ -358,15 +369,15 @@ def _alloc_pointed_array_on_high_bw_mem(self, site, obj, storage):
 
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
-        nbytes = SizeOf(obj._C_typedata, stars='*')*obj.dim.symbolic_size
+        nbytes = SizeOf(obj._C_typedata, stars='*')*self.intm(obj.dim.symbolic_size)
         alloc0 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         free0 = self.langbb['host-free'](obj._C_symbol)
 
         # The pointee Array
         pobj = IndexedPointer(obj._C_symbol, obj.dim)
         memptr = VOID(Byref(pobj), '**')
-        nbytes = SizeOf(obj._C_typedata)*obj.array.size
+        nbytes = SizeOf(obj._C_typedata)*self.intm(obj.array.size)
         alloc1 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         free1 = self.langbb['host-free'](pobj)

devito/symbolics/extended_dtypes.py

@@ -7,7 +7,8 @@
 from devito.tools.dtypes_lowering import dtype_mapper
 
 __all__ = ['cast', 'CustomType', 'limits_mapper', 'INT', 'FLOAT', 'BaseCast',  # noqa
-           'DOUBLE', 'VOID', 'NoDeclStruct', 'c_complex', 'c_double_complex']
+           'DOUBLE', 'VOID', 'NoDeclStruct', 'c_complex', 'c_double_complex',
+           'LONG']
 
 
 limits_mapper = {
@@ -72,6 +73,7 @@ def cast(casttype, stars=None):
 
 ULONG = cast(np.uint64)
 UINTP = cast(np.uint32, '*')
+LONG = cast(np.int64)
 
 
 # Standard ones, needed as class for e.g. single dispatch

tests/test_iet.py

@@ -345,7 +345,7 @@ class SpecialObject2(LocalObject):
     assert 'bar<int,float>& obj2;' in str(iet)
     assert 'dummy obj3 = meh;' in str(iet)
     assert 'dummy obj4(1,2) = meh;' in str(iet)
-    assert 'dummy obj5(ceil(1.0F/(float)s)) = meh;' in str(iet)
+    assert 'dummy obj5(ceil(1.0F/(float)(s))) = meh;' in str(iet)
     assert 'float obj6 = meh;' in str(iet)
 
 

tests/test_linearize.py

@@ -3,7 +3,7 @@
 import scipy.sparse
 
 from devito import (Grid, Function, TimeFunction, SparseTimeFunction, Operator, Eq,
-                    Inc, MatrixSparseTimeFunction, sin, switchconfig)
+                    Inc, MatrixSparseTimeFunction, sin, switchconfig, configuration)
 from devito.ir import Call, Callable, DummyExpr, Expression, FindNodes, SymbolRegistry
 from devito.passes import Graph, linearize, generate_macros
 from devito.types import Array, Bundle, DefaultDimension
@@ -640,3 +640,18 @@ def _test_different_dtype():
         assert "L0(x,y) f[(x)*y_stride0 + (y)]" in str(op1)
 
     _test_different_dtype()
+
+
+@pytest.mark.parametrize('order', [2, 4])
+def test_int64_array(order):
+
+    grid = Grid(shape=(4, 4))
+    f = Function(name='f', grid=grid, space_order=order)
+
+    a = Array(name='a', dimensions=grid.dimensions, shape=grid.shape,
+              halo=f.halo)
+
+    eqs = [Eq(f, a.indexify() + 1)]
+    op = Operator(eqs, opt=('advanced', {'linearize': True, 'index-mode': 'int64'}))
+    long = 'static_cast<long>' if 'CXX' in configuration['language'] else '(long)'
+    assert f'({2*order} + {long}(y_size))*({2*order} + {long}(x_size)))' in str(op)

tests/test_symbolics.py

@@ -426,10 +426,10 @@ class BarCast(BaseCast):
         _dtype = 'bar'
 
     v = BarCast(s, '**')
-    assert ccode(v) == '(bar**)s'
+    assert ccode(v) == '(bar**)(s)'
 
     # Reconstruction
-    assert ccode(v.func(*v.args)) == '(bar**)s'
+    assert ccode(v.func(*v.args)) == '(bar**)(s)'
 
     v1 = BarCast(s, '****')
     assert v != v1
@@ -442,19 +442,19 @@ def test_str_cast():
     assert not v.stars
     assert v.dtype == 'foo'
     assert v._op == '(foo)'
-    assert ccode(v) == '(foo)s'
+    assert ccode(v) == '(foo)(s)'
 
     v = Cast(s, 'foo*')
     assert v.stars == '*'
     assert v.dtype == 'foo'
     assert v._op == '(foo*)'
-    assert ccode(v) == '(foo*)s'
+    assert ccode(v) == '(foo*)(s)'
 
     v = Cast(s, 'foo **')
     assert v.stars == '**'
     assert v.dtype == 'foo'
     assert v._op == '(foo**)'
-    assert ccode(v) == '(foo**)s'
+    assert ccode(v) == '(foo**)(s)'
 
 
 def test_findexed():