compiler: Revamp Array allocation

Author: FabioLuporini

devito/operator/operator.py

@@ -13,11 +13,12 @@
 from devito.data import default_allocator
 from devito.exceptions import (CompilationError, ExecutionError, InvalidArgument,
                                InvalidOperator)
-from devito.logger import debug, info, perf, warning, is_log_enabled_for, switch_log_level
+from devito.logger import (debug, info, perf, warning, is_log_enabled_for,
+                           switch_log_level)
 from devito.ir.equations import LoweredEq, lower_exprs, concretize_subdims
 from devito.ir.clusters import ClusterGroup, clusterize
-from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols, MetaCall,
-                           derive_parameters, iet_build)
+from devito.ir.iet import (Callable, CInterface, EntryFunction, FindSymbols,
+                           MetaCall, derive_parameters, iet_build)
 from devito.ir.support import AccessMode, SymbolRegistry
 from devito.ir.stree import stree_build
 from devito.operator.profiling import create_profile
@@ -26,8 +27,7 @@
 from devito.parameters import configuration
 from devito.passes import (Graph, lower_index_derivatives, generate_implicit,
                            generate_macros, minimize_symbols, unevaluate,
-                           error_mapper, is_on_device)
-from devito.passes.iet.dtypes import lower_dtypes
+                           error_mapper, is_on_device, lower_dtypes)
 from devito.symbolics import estimate_cost, subs_op_args
 from devito.tools import (DAG, OrderedSet, Signer, ReducerMap, as_mapper, as_tuple,
                           flatten, filter_sorted, frozendict, is_integer,
@@ -488,7 +488,7 @@ def _lower_iet(cls, uiet, profiler=None, **kwargs):
         # Extract the necessary macros from the symbolic objects
         generate_macros(graph, **kwargs)
 
-        # Add type specific metadata
+        # Target-specific lowering
         lower_dtypes(graph, **kwargs)
 
         # Target-independent optimizations

devito/passes/iet/definitions.py

@@ -9,17 +9,22 @@
 
 import numpy as np
 
-from devito.ir import (Block, Call, Definition, DummyExpr, Return, EntryFunction,
-                       FindNodes, FindSymbols, MapExprStmts, Transformer,
-                       make_callable)
+from devito.ir import (
+    Block, Call, Definition, DummyExpr, Iteration, List, Return, EntryFunction,
+    FindNodes, FindSymbols, MapExprStmts, Transformer, make_callable
+)
 from devito.passes import is_gpu_create
 from devito.passes.iet.engine import iet_pass
 from devito.passes.iet.langbase import LangBB
-from devito.symbolics import (Byref, DefFunction, FieldFromPointer, IndexedPointer,
-                              SizeOf, VOID, pow_to_mul, unevaluate)
+from devito.symbolics import (
+    Byref, DefFunction, FieldFromPointer, IndexedPointer, ListInitializer,
+    SizeOf, VOID, pow_to_mul, unevaluate
+)
 from devito.tools import as_mapper, as_list, as_tuple, filter_sorted, flatten
-from devito.types import (Array, ComponentAccess, CustomDimension, DeviceMap,
-                          DeviceRM, Eq, Symbol, size_t)
+from devito.types import (
+    Array, ComponentAccess, CustomDimension, Dimension, DeviceMap, DeviceRM,
+    Eq, Symbol, size_t
+)
 
 __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage']
 
@@ -40,6 +45,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
         self.defined = set()
+        self.includes = set()
 
     def update(self, key, site, **kwargs):
         if key in self.defined:
@@ -67,6 +73,10 @@ def map(self, key, site, k, v):
 
         self.defined.add((site[-1], key))
 
+    def include(self, v):
+        if v:
+            self.includes.add(v)
+
 
 class DataManager:
 
@@ -132,8 +142,7 @@ def _alloc_array_on_global_mem(self, site, obj, storage):
         alloc = Call(name, efunc.parameters)
 
         storage.update(obj, site, allocs=alloc, efuncs=efunc)
-
-        return self.langbb['header-memcpy']
+        storage.include(self.langbb['header-memcpy'])
 
     def _alloc_scalar_on_low_lat_mem(self, site, expr, storage):
         """
@@ -170,6 +179,12 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
 
         arity_param = Symbol(name='arity', dtype=size_t)
         arity_arg = SizeOf(obj.indexed._C_typedata)
+        ndims_param = Symbol(name='ndims', dtype=size_t)
+        ndims_arg = obj.ndim
+        shape_param = Array(name=f'{obj.name}_shape', dtype=np.uint64,
+                            dimensions=(Dimension(name='d'),), scope='rvalue')
+        shape_arg = ListInitializer(obj.c0.symbolic_shape, dtype=shape_param.dtype)
+
         ffp0 = FieldFromPointer(obj._C_field_data, obj._C_symbol)
         ffp1 = FieldFromPointer(obj._C_field_shape, obj._C_symbol)
         ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
@@ -183,14 +198,21 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
 
         # Allocate the shape array
         memptr = VOID(Byref(ffp1), '**')
-        nbytes = SizeOf(obj._C_size_type)*obj.ndim
+        nbytes = SizeOf(obj._C_size_type)*ndims_param
         alloc1 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
-        # Initialize the Array struct
-        init = [*[DummyExpr(IndexedPointer(ffp1, i), s)
-                  for i, s in enumerate(obj.c0.symbolic_shape)],
-                DummyExpr(ffp2, obj.size),
-                DummyExpr(ffp3, ffp2*arity_param)]
+        # Initialize the Array metadata
+        dim, = shape_param.dimensions
+        init = [DummyExpr(ffp2, 1)]
+        init.append(Iteration(
+            List(body=(
+                DummyExpr(IndexedPointer(ffp1, dim), shape_param[dim]),
+                DummyExpr(ffp2, ffp2*shape_param[dim])
+            )),
+            dim,
+            ndims_param - 1,
+        ))
+        init.append(DummyExpr(ffp3, ffp2*arity_param))
 
         # Allocate the underlying host data
         memptr = VOID(Byref(ffp0), '**')
@@ -213,6 +235,8 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         efunc0 = make_callable(name, body, retval=obj)
         args = list(efunc0.parameters)
         args[args.index(arity_param)] = arity_arg
+        args[args.index(ndims_param)] = ndims_arg
+        args[args.index(shape_param)] = shape_arg
         alloc = Call(name, args, retobj=obj)
 
         # Same story for the frees
@@ -222,6 +246,7 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         free = Call(name, efunc1.parameters)
 
         storage.update(obj, site, allocs=alloc, frees=free, efuncs=(efunc0, efunc1))
+        storage.include(self.langbb['header-array'])
 
     def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         """
@@ -436,18 +461,15 @@ def place_definitions(self, iet, globs=None, **kwargs):
                 self._alloc_pointed_array_on_high_bw_mem(iet, i, storage)
 
         # Handle postponed global objects
-        includes = set()
         if isinstance(iet, EntryFunction) and globs:
             for i in sorted(globs, key=lambda f: f.name):
-                v = self._alloc_array_on_global_mem(iet, i, storage)
-                if v:
-                    includes.add(v)
+                self._alloc_array_on_global_mem(iet, i, storage)
 
         iet, efuncs = self._inject_definitions(iet, storage)
 
         return iet, {'efuncs': efuncs,
                      'globals': as_tuple(globs),
-                     'includes': as_tuple(includes)}
+                     'includes': as_tuple(sorted(storage.includes))}
 
     @iet_pass
     def place_casts(self, iet, **kwargs):

devito/passes/iet/languages/C.py

@@ -5,14 +5,17 @@
 from devito.passes.iet.definitions import DataManager
 from devito.passes.iet.orchestration import Orchestrator
 from devito.passes.iet.langbase import LangBB
-from devito.symbolics.extended_dtypes import c_complex, c_double_complex
+from devito.symbolics import c_complex, c_double_complex
+from devito.tools import dtype_to_cstr
 
 __all__ = ['CBB', 'CDataManager', 'COrchestrator']
 
 
 class CBB(LangBB):
 
     mapper = {
+        # Misc
+        'header-array': None,
         # Complex
         'includes-complex': 'complex.h',
         # Allocs
@@ -55,3 +58,12 @@ class CPrinter(BasePrinter, C99CodePrinter):
 
     def _print_ImaginaryUnit(self, expr):
         return '_Complex_I'
+
+    def _print_ListInitializer(self, expr):
+        li = super()._print_ListInitializer(expr)
+        if expr.dtype:
+            # C99, unlike CXX, supports compound literals
+            tstr = dtype_to_cstr(expr.dtype)
+            return f'({tstr}[]){li}'
+        else:
+            return li

devito/passes/iet/languages/CXX.py

@@ -3,7 +3,8 @@
 
 from devito.ir import Call, UsingNamespace, BasePrinter
 from devito.passes.iet.langbase import LangBB
-from devito.symbolics.extended_dtypes import c_complex, c_double_complex
+from devito.symbolics import c_complex, c_double_complex
+from devito.tools import dtype_to_cstr
 
 __all__ = ['CXXBB']
 
@@ -65,6 +66,8 @@ def std_arith(prefix=None):
 class CXXBB(LangBB):
 
     mapper = {
+        # Misc
+        'header-array': 'array',
         # Complex
         'includes-complex': 'complex',
         'complex-namespace': [UsingNamespace('std::complex_literals')],
@@ -112,3 +115,12 @@ def _print_Cast(self, expr):
         caster = 'reinterpret_cast' if expr.reinterpret else 'static_cast'
         cast = f'{caster}<{tstr}{self._print(expr.stars)}>'
         return self._print_UnaryOp(expr, op=cast, parenthesize=True)
+
+    def _print_ListInitializer(self, expr):
+        li = super()._print_ListInitializer(expr)
+        if expr.dtype:
+            # CXX, unlike C99, does not support compound literals
+            tstr = dtype_to_cstr(expr.dtype)
+            return f'std::array<{tstr}, {len(expr.params)}>{li}.data()'
+        else:
+            return li

devito/symbolics/extended_sympy.py

@@ -292,16 +292,20 @@ class ListInitializer(sympy.Expr, Pickable):
     """
 
     __rargs__ = ('params',)
+    __rkwargs__ = ('dtype',)
 
-    def __new__(cls, params):
+    def __new__(cls, params, dtype=None):
         args = []
         for p in as_tuple(params):
             try:
                 args.append(sympify(p))
             except sympy.SympifyError:
                 raise ValueError(f"Illegal param `{p}`")
         obj = sympy.Expr.__new__(cls, *args)
+
         obj.params = tuple(args)
+        obj.dtype = dtype
+
         return obj
 
     def __str__(self):

devito/types/array.py

@@ -40,7 +40,7 @@ def __indices_setup__(cls, *args, **kwargs):
 
     @property
     def _C_name(self):
-        if self._mem_stack or self._mem_constant:
+        if self._mem_stack or self._mem_constant or self._mem_rvalue:
             # No reason to distinguish between two different names, that is
             # the _C_name and the name -- just `self.name` is enough
             return self.name
@@ -97,7 +97,7 @@ class Array(ArrayBasic):
         to 'local'. Used to override `_mem_local` and `_mem_mapped`.
     scope : str, optional
         The scope in the given memory space. Allowed values: 'heap', 'stack',
-        'static', 'constant', 'shared', 'shared-remote', 'registers'.
+        'static', 'constant', 'shared', 'shared-remote', 'registers', 'rvalue'.
         'static' refers to a static array in a C/C++ sense. 'constant' and
         'shared' mean that the Array represents an object allocated in so
         called constant and shared memory, respectively, which are typical of
@@ -107,7 +107,9 @@ class Array(ArrayBasic):
         architecture doesn't have something akin to constant memory, the Array
         falls back to a global, const, static array in a C/C++ sense.
         'registers' is used to indicate that the Array has a small static size
-        and, as such, it could be allocated in registers. Defaults to 'heap'.
+        and, as such, it could be allocated in registers. If `rvalue`, the
+        Array is treated as a temporary or "transient" object, just like
+        C++'s rvalue references or C's compound literals. Defaults to 'heap'.
         Note that not all scopes make sense for a given space.
     grid : Grid, optional
         Only necessary for distributed-memory parallelism; a Grid contains
@@ -142,7 +144,7 @@ def __init_finalize__(self, *args, **kwargs):
 
         self._scope = kwargs.get('scope', 'heap')
         assert self._scope in ['heap', 'stack', 'static', 'constant', 'shared',
-                               'shared-remote', 'registers']
+                               'shared-remote', 'registers', 'rvalue']
 
         self._initvalue = kwargs.get('initvalue')
         assert self._initvalue is None or self._scope != 'heap'
@@ -197,6 +199,10 @@ def _mem_registers(self):
     def _mem_constant(self):
         return self._scope == 'constant'
 
+    @property
+    def _mem_rvalue(self):
+        return self._scope == 'rvalue'
+
     @property
     def initvalue(self):
         return self._initvalue
@@ -473,10 +479,11 @@ def initvalue(self):
     # Defaulting to self.c0's behaviour
     for i in ('_mem_internal_eager', '_mem_internal_lazy', '_mem_local',
               '_mem_mapped', '_mem_host', '_mem_stack', '_mem_constant',
-              '_mem_shared', '_mem_shared_remote', '__padding_dtype__',
-              '_size_domain', '_size_halo', '_size_owned', '_size_padding',
-              '_size_nopad', '_size_nodomain', '_offset_domain', '_offset_halo',
-              '_offset_owned', '_dist_dimensions', '_C_get_field', 'grid',
+              '_mem_shared', '_mem_shared_remote', '_mem_registers',
+              '_mem_rvalue', '__padding_dtype__', '_size_domain', '_size_halo',
+              '_size_owned', '_size_padding', '_size_nopad', '_size_nodomain',
+              '_offset_domain', '_offset_halo', '_offset_owned',
+              '_dist_dimensions', '_C_get_field', 'grid',
               *AbstractFunction.__properties__):
         locals()[i] = property(lambda self, v=i: getattr(self.c0, v))
 
@@ -517,7 +524,6 @@ def _C_ctype(self):
         if self._mem_mapped:
             return super()._C_ctype
         else:
-            #TODO DROP???
             return POINTER(dtype_to_ctype(self.dtype))
 
 

devito/types/basic.py

@@ -42,7 +42,8 @@ class CodeSymbol:
         * "liveness": `_mem_external`, `_mem_internal_eager`, `_mem_internal_lazy`
         * "space": `_mem_local`, `_mem_mapped`, `_mem_host`
         * "scope": `_mem_stack`, `_mem_heap`, `_mem_global`, `_mem_shared`,
-                   `_mem_shared_remote`, `_mem_constant`
+                   `_mem_shared_remote`, `_mem_constant`, `_mem_registers`,
+                   `_mem_rvalue`
 
     For example, an object that is `<_mem_internal_lazy, _mem_local, _mem_heap>`
     is allocated within the Operator entry point, on either the host or device
@@ -230,6 +231,21 @@ def _mem_shared_remote(self):
         """
         return False
 
+    @property
+    def _mem_registers(self):
+        """
+        True if the associated data is allocated in registers, False otherwise.
+        """
+        return False
+
+    @property
+    def _mem_rvalue(self):
+        """
+        True if the associated data is allocated in a temporary (or "transient")
+        variable, such as rvalues in CXX, False otherwise.
+        """
+        return False
+
 
 class Basic(CodeSymbol):