moar moar

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -229,35 +229,50 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         """
         decl = Definition(obj)
 
+        arity_param = Symbol(name='arity', dtype=size_t)
+        arity_arg = SizeOf(obj.indexed._C_typedata)
+        ffp1 = FieldFromPointer(obj._C_field_shape, obj._C_symbol)
+        ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
+        ffp3 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
+
         # Allocate the Bundle struct
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
         nbytes = SizeOf(obj._C_typedata)
-        alloc = self.langbb['host-alloc'](memptr, alignment, nbytes)
+        alloc0 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
-        nbytes_param = Symbol(name='nbytes', dtype=np.uint64, is_const=True)
-        nbytes_arg = SizeOf(obj.indexed._C_typedata)*obj.size
+        # Allocate the shape array
+        memptr = VOID(Byref(ffp1), '**')
+        nbytes = SizeOf(obj._C_size_type)*obj.ndim
+        alloc1 = self.langbb['host-alloc'](memptr, alignment, nbytes)
 
         # Initialize the Bundle struct
-        ffp1 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
-        init0 = DummyExpr(ffp1, nbytes_param)
-        ffp2 = FieldFromPointer(obj._C_field_size, obj._C_symbol)
-        init1 = DummyExpr(ffp2, 0)
+        init = [*[DummyExpr(IndexedPointer(ffp1, i), s)
+                  for i, s in enumerate(obj.c0.symbolic_shape)],
+                DummyExpr(ffp2, obj.size),
+                DummyExpr(ffp3, ffp2*arity_param)]
 
-        free = self.langbb['host-free'](obj._C_symbol)
+        # Free all of the allocated data
+        frees = [self.langbb['host-free'](ffp1),
+                 self.langbb['host-free'](obj._C_symbol)]
 
         ret = Return(obj._C_symbol)
 
         # Wrap everything in a Callable so that we can reuse the same code
         # for equivalent Bundle structs
         name = self.sregistry.make_name(prefix='alloc')
-        body = (decl, alloc, init0, init1, ret)
+        body = (decl, alloc0, alloc1, *init, ret)
         efunc0 = make_callable(name, body, retval=obj)
         args = list(efunc0.parameters)
-        args[args.index(nbytes_param)] = nbytes_arg
+        args[args.index(arity_param)] = arity_arg
         alloc = Call(name, args, retobj=obj)
 
-        storage.update(obj, site, allocs=alloc, frees=free, efuncs=efunc0)
+        # Same story for the frees
+        name = self.sregistry.make_name(prefix='free')
+        efunc1 = make_callable(name, frees)
+        free = Call(name, efunc1.parameters)
+
+        storage.update(obj, site, allocs=alloc, frees=free, efuncs=(efunc0, efunc1))
 
     def _alloc_object_array_on_low_lat_mem(self, site, obj, storage):
         """

devito/types/array.py

@@ -4,7 +4,6 @@
 import numpy as np
 from sympy import Expr, cacheit
 
-from devito.data import FULL
 from devito.tools import (Pickable, as_tuple, c_restrict_void_p,
                           dtype_to_ctype, dtypes_vector_mapper, is_integer)
 from devito.types.basic import AbstractFunction, LocalType