Merge pull request #2664 from devitocodes/index-mode

compiler:  fix temp/array alloc

Author: mloubout

devito/arch/compiler.py

@@ -468,7 +468,7 @@ def __init_finalize__(self, **kwargs):
                 # they support `#pragma omp simd`
                 self.ldflags += ['-fopenmp']
         except (TypeError, ValueError):
-            if language == 'openmp':
+            if 'openmp' in language:
                 self.ldflags += ['-fopenmp']
 
     def __lookup_cmds__(self):
@@ -517,7 +517,7 @@ def __init_finalize__(self, **kwargs):
             # The compiler can be installed with Homebrew or can be built from scratch.
             # Check if installed and set compiler flags accordingly
             llvmm1 = get_m1_llvm_path(language)
-            if llvmm1 and language == 'openmp':
+            if llvmm1 and 'openmp' in language:
                 mx = platform.march
                 self.ldflags += [f'-mcpu=apple-{mx}',
                                  '-fopenmp', f'-L{llvmm1["libs"]}']
@@ -528,7 +528,7 @@ def __init_finalize__(self, **kwargs):
                 self.cflags += ['-mcpu=native']
             else:
                 self.cflags += ['-march=native']
-            if language == 'openmp':
+            if 'openmp' in language:
                 self.ldflags += ['-fopenmp']
 
     def __lookup_cmds__(self):
@@ -562,7 +562,7 @@ def __init_finalize__(self, **kwargs):
         if not configuration['safe-math']:
             self.cflags.append('-ffast-math')
 
-        if language == 'openmp':
+        if 'openmp' in language:
             self.ldflags += ['-fopenmp']
 
         if isinstance(platform, NvidiaDevice):
@@ -621,10 +621,10 @@ def __init_finalize__(self, **kwargs):
 
             if language == 'openacc':
                 self.cflags.extend(['-mp', '-acc:gpu'])
-            elif language == 'openmp':
+            elif 'openmp' in language:
                 self.cflags.extend(['-mp=gpu'])
         elif isinstance(platform, Cpu64):
-            if language == 'openmp':
+            if 'openmp' in language:
                 self.cflags.append('-mp')
             if isinstance(platform, NvidiaArm):
                 self.cflags.append(f'-mcpu={platform.march}')
@@ -770,7 +770,7 @@ def __init_finalize__(self, **kwargs):
             # Systematically use 512-bit vectors if avx512 is available.
             self.cflags.append("-qopt-zmm-usage=high")
 
-        if language == 'openmp':
+        if 'openmp' in language:
             self.ldflags.append('-qopenmp')
 
         if kwargs.get('mpi'):
@@ -841,7 +841,7 @@ def __init_finalize__(self, **kwargs):
         if language == 'sycl':
             warning(f"Use SyclCompiler (`sycl`) to jit-compile sycl, not {self.name}")
 
-        elif language == 'openmp':
+        elif 'openmp' in language:
             # Earlier versions to OneAPI 2023.2.0 (clang17 underneath), have an
             # OpenMP bug concerning reductions, hence with them we're forced to
             # use the obsolete -fopenmp

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=not expr.base.is_Atom)
 
     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, as_long
 )
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
 from devito.types import (
@@ -136,7 +136,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)*as_long(obj.size)
         body = [Definition(src),
                 self.langbb['alloc-global-symbol'](obj.indexed, src.indexed, nbytes)]
         efunc = make_callable(name, body)
@@ -159,7 +159,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)*as_long(obj.size)
         alloc = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         free = self.langbb['host-free'](obj._C_symbol)
@@ -358,15 +358,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='*')*as_long(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)*as_long(obj.array.size)
         alloc1 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         free1 = self.langbb['host-free'](pobj)
@@ -551,10 +551,10 @@ def process(self, graph):
 
 class DeviceAwareDataManager(DataManager):
 
-    def __init__(self, **kwargs):
-        self.gpu_fit = kwargs['options']['gpu-fit']
-        self.gpu_create = kwargs['options']['gpu-create']
-        self.pmode = kwargs['options'].get('place-transfers')
+    def __init__(self, options=None, **kwargs):
+        self.gpu_fit = options['gpu-fit']
+        self.gpu_create = options['gpu-create']
+        self.pmode = options.get('place-transfers')
 
         super().__init__(**kwargs)
 

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

devito/symbolics/manipulation.py

@@ -11,8 +11,11 @@
     EvalDerivative, IndexDerivative
 )
 from devito.symbolics.extended_sympy import DefFunction, rfunc
+from devito.symbolics.extended_dtypes import LONG
 from devito.symbolics.queries import q_leaf
-from devito.symbolics.search import retrieve_indexed, retrieve_functions
+from devito.symbolics.search import (
+    retrieve_indexed, retrieve_functions, retrieve_symbols
+)
 from devito.symbolics.unevaluation import (
     Add as UnevalAdd, Mul as UnevalMul, Pow as UnevalPow, UnevaluableMixin
 )
@@ -24,7 +27,8 @@
 
 __all__ = ['xreplace_indices', 'pow_to_mul', 'indexify', 'subs_op_args',
            'normalize_args', 'uxreplace', 'Uxmapper', 'subs_if_composite',
-           'reuse_if_untouched', 'evalrel', 'flatten_args', 'unevaluate']
+           'reuse_if_untouched', 'evalrel', 'flatten_args', 'unevaluate',
+           'as_long']
 
 
 def uxreplace(expr, rule):
@@ -523,3 +527,14 @@ def unevaluate(expr):
         return uneval_mapper[expr.func](*args)
     except KeyError:
         return reuse_if_untouched(expr, args)
+
+
+def as_long(expr):
+    """
+    Convert an expression and its symbolic args to a long integer.
+    """
+    try:
+        syms = retrieve_symbols(expr)
+        return expr.subs({s: LONG(s) for s in syms})
+    except AttributeError:
+        return LONG(expr)

examples/mpi/overview.ipynb

@@ -486,9 +486,9 @@
        "  MPI_Request rsend;\n",
        "\n",
        "  float *restrict bufg_vec __attribute__ ((aligned (64)));\n",
-       "  posix_memalign((void**)(&bufg_vec),64,x_size*y_size*sizeof(float));\n",
+       "  posix_memalign((void**)(&bufg_vec),64,sizeof(float)*(long)y_size*(long)x_size);\n",
        "  float *restrict bufs_vec __attribute__ ((aligned (64)));\n",
-       "  posix_memalign((void**)(&bufs_vec),64,x_size*y_size*sizeof(float));\n",
+       "  posix_memalign((void**)(&bufs_vec),64,sizeof(float)*(long)y_size*(long)x_size);\n",
        "\n",
        "  MPI_Irecv(bufs_vec,x_size*y_size,MPI_FLOAT,fromrank,13,comm,&(rrecv));\n",
        "  if (torank != MPI_PROC_NULL)\n",

examples/performance/00_overview.ipynb

@@ -1192,13 +1192,13 @@
       "int Kernel(struct dataobj *restrict f_vec, struct dataobj *restrict u_vec, const float h_y, const int time_M, const int time_m, const int x0_blk0_size, const int x_M, const int x_m, const int y0_blk0_size, const int y_M, const int y_m, const int z_M, const int z_m, const int nthreads, const int z_size, const int x_size, const int y_size, struct profiler * timers)\n",
       "{\n",
       "  float **restrict pr2_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&pr2_vec),64,nthreads*sizeof(float*));\n",
+      "  posix_memalign((void**)(&pr2_vec),64,sizeof(float*)*(long)nthreads);\n",
       "  float *restrict r0_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&r0_vec),64,x_size*y_size*z_size*sizeof(float));\n",
+      "  posix_memalign((void**)(&r0_vec),64,sizeof(float)*(long)z_size*(long)y_size*(long)x_size);\n",
       "  #pragma omp parallel num_threads(nthreads)\n",
       "  {\n",
       "    const int tid = omp_get_thread_num();\n",
-      "    posix_memalign((void**)(&(pr2_vec[tid])),64,z_size*(y0_blk0_size + 4)*sizeof(float));\n",
+      "    posix_memalign((void**)(&(pr2_vec[tid])),64,sizeof(float)*(long)z_size*(4 + (long)y0_blk0_size));\n",
       "  }\n",
       "\n",
       "  float (*restrict f)[f_vec->size[1]][f_vec->size[2]] __attribute__ ((aligned (64))) = (float (*)[f_vec->size[1]][f_vec->size[2]]) f_vec->data;\n",
@@ -1412,7 +1412,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "posix_memalign((void**)(&r0_vec),64,x_size*y_size*z_size*sizeof(float));\n"
+      "posix_memalign((void**)(&r0_vec),64,sizeof(float)*(long)z_size*(long)y_size*(long)x_size);\n"
      ]
     }
    ],
@@ -1483,13 +1483,13 @@
       "int Kernel(struct dataobj *restrict f_vec, struct dataobj *restrict u_vec, const float h_y, const int time_M, const int time_m, const int x_M, const int x_m, const int y_M, const int y_m, const int z_M, const int z_m, const int nthreads, const int z_size, const int x_size, const int y_size, struct profiler * timers)\n",
       "{\n",
       "  float **restrict pr2_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&pr2_vec),64,nthreads*sizeof(float*));\n",
+      "  posix_memalign((void**)(&pr2_vec),64,sizeof(float*)*(long)nthreads);\n",
       "  float *restrict r0_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&r0_vec),64,x_size*y_size*z_size*sizeof(float));\n",
+      "  posix_memalign((void**)(&r0_vec),64,sizeof(float)*(long)z_size*(long)y_size*(long)x_size);\n",
       "  #pragma omp parallel num_threads(nthreads)\n",
       "  {\n",
       "    const int tid = omp_get_thread_num();\n",
-      "    posix_memalign((void**)(&(pr2_vec[tid])),64,z_size*(y_size + 4)*sizeof(float));\n",
+      "    posix_memalign((void**)(&(pr2_vec[tid])),64,sizeof(float)*(long)z_size*(4 + (long)y_size));\n",
       "  }\n",
       "\n",
       "  float (*restrict f)[f_vec->size[1]][f_vec->size[2]] __attribute__ ((aligned (64))) = (float (*)[f_vec->size[1]][f_vec->size[2]]) f_vec->data;\n",
@@ -1626,11 +1626,11 @@
       "int Kernel(struct dataobj *restrict f_vec, struct dataobj *restrict u_vec, const float h_x, const float h_y, const int time_M, const int time_m, const int x0_blk0_size, const int x1_blk0_size, const int x_M, const int x_m, const int y0_blk0_size, const int y1_blk0_size, const int y_M, const int y_m, const int z_M, const int z_m, const int nthreads, const int x_size, const int y_size, const int z_size, struct profiler * timers)\n",
       "{\n",
       "  float *restrict r0_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&r0_vec),64,x_size*y_size*z_size*sizeof(float));\n",
+      "  posix_memalign((void**)(&r0_vec),64,sizeof(float)*(long)z_size*(long)y_size*(long)x_size);\n",
       "  float *restrict r3_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&r3_vec),64,z_size*(x_size + 4)*(y_size + 4)*sizeof(float));\n",
+      "  posix_memalign((void**)(&r3_vec),64,sizeof(float)*(long)z_size*(4 + (long)y_size)*(4 + (long)x_size));\n",
       "  float *restrict r4_vec __attribute__ ((aligned (64)));\n",
-      "  posix_memalign((void**)(&r4_vec),64,z_size*(x_size + 4)*(y_size + 4)*sizeof(float));\n",
+      "  posix_memalign((void**)(&r4_vec),64,sizeof(float)*(long)z_size*(4 + (long)y_size)*(4 + (long)x_size));\n",
       "\n",
       "  float (*restrict f)[f_vec->size[1]][f_vec->size[2]] __attribute__ ((aligned (64))) = (float (*)[f_vec->size[1]][f_vec->size[2]]) f_vec->data;\n",
       "  float (*restrict r0)[y_size][z_size] __attribute__ ((aligned (64))) = (float (*)[y_size][z_size]) r0_vec;\n",

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,22 @@ 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'}))
+    if 'CXX' in configuration['language']:
+        long = 'static_cast<long>'
+        assert f'({2*order} + {long}(y_size))*({2*order} + {long}(x_size)))' in str(op)
+    else:
+        long = '(long)'
+        assert f'({2*order} + {long}y_size)*({2*order} + {long}x_size))' in str(op)