WIP WIP

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -177,21 +177,23 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
         decl = Definition(obj)
 
+        # NOTE: the `arity` is calculated such as `sizeof(float3)/sizeof(float)`
+        # for portability reasons (since we don't know the size of compound
+        # types a priori)
         arity_param = Symbol(name='arity', dtype=size_t)
-        arity_arg = SizeOf(obj.indexed._C_typedata)
+        arity_arg = (SizeOf(obj.indexed._C_typedata) /
+                     SizeOf(obj.c0.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)
-        sizeofelem_param = Symbol(name='sizeofelem', dtype=size_t)
-        sizeofelem_arg = 1  #TODO
 
         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)
         ffp3 = FieldFromPointer(obj._C_field_nbytes, obj._C_symbol)
-        ffp4 = FieldFromPointer(obj._C_field_sizeofelem, obj._C_symbol)
+        ffp4 = FieldFromPointer(obj._C_field_arity, obj._C_symbol)
 
         # Allocate the Array struct
         memptr = VOID(Byref(obj._C_symbol), '**')
@@ -216,7 +218,7 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
             ndims_param - 1,
         ))
         init.append(DummyExpr(ffp3, ffp2*arity_param))
-        init.append(DummyExpr(ffp4, sizeofelem_param))
+        init.append(DummyExpr(ffp4, arity_param))
 
         # Allocate the underlying host data
         memptr = VOID(Byref(ffp0), '**')
@@ -241,7 +243,6 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         args[args.index(arity_param)] = arity_arg
         args[args.index(ndims_param)] = ndims_arg
         args[args.index(shape_param)] = shape_arg
-        args[args.index(sizeofelem_param)] = sizeofelem_arg  #TODO
         alloc = Call(name, args, retobj=obj)
 
         # Same story for the frees
@@ -259,21 +260,22 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         """
         decl = Definition(obj)
 
+        # NOTE: the `arity` is calculated such as `sizeof(float3)/sizeof(float)`
+        # for portability reasons (since we don't know the size of compound
+        # types a priori)
         arity_param = Symbol(name='arity', dtype=size_t)
-        arity_arg = SizeOf(obj.indexed._C_typedata)
+        arity_arg = (SizeOf(obj.indexed._C_typedata) /
+                     SizeOf(obj.c0.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)
-        sizeofelem_param = Symbol(name='sizeofelem', dtype=size_t)
-        sizeofelem_arg = 1  #TODO
-        from IPython import embed; embed()
 
         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)
-        ffp4 = FieldFromPointer(obj._C_field_sizeofelem, obj._C_symbol)
+        ffp4 = FieldFromPointer(obj._C_field_arity, obj._C_symbol)
 
         # Allocate the Bundle struct
         memptr = VOID(Byref(obj._C_symbol), '**')
@@ -298,6 +300,7 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
             ndims_param - 1,
         ))
         init.append(DummyExpr(ffp3, ffp2*arity_param))
+        init.append(DummyExpr(ffp4, arity_param))
 
         # Free all of the allocated data
         frees = [self.langbb['host-free'](ffp1),
@@ -314,7 +317,6 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         args[args.index(arity_param)] = arity_arg
         args[args.index(ndims_param)] = ndims_arg
         args[args.index(shape_param)] = shape_arg
-        args[args.index(sizeofelem_param)] = sizeofelem_arg  #TODO
         alloc = Call(name, args, retobj=obj)
 
         # Same story for the frees

devito/passes/iet/engine.py

@@ -361,24 +361,30 @@ def abstract_component_accesses(root, efuncs, sregistry=None):
         ca, = compaccs
         f = ca.function
 
-        # Access the same entry as `ca` but via pointer arithmetic instead
+        arity = Symbol(name='arity', dtype=np.uint32, is_const=True)
         index = Symbol(name='index', dtype=np.uint32, is_const=True)
-        indices = [Mul(f.ncomp, i, evaluate=False) for i in ca.indices]  #TODO
+
+        # Access the same entry as `ca` but via pointer arithmetic instead
+        indices = [Mul(arity, i, evaluate=False) for i in ca.indices]
         indices[-1] += index
         kernel1 = Uxreplace({ca: f.c0.indexed[indices]}).visit(kernel)
 
         # Update the parameters list
-        parameters = [*kernel1.parameters, index]
+        parameters = [*kernel1.parameters, index, arity]
         parameters[parameters.index(f.indexed)] = f.c0.indexed
         kernel1 = kernel1._rebuild(parameters=parameters)
 
         # Update the Call site
-        args = [*call.arguments, ca.index]
-        args[args.index(f.dmap)] = Cast(f.dmap, dtype=f.c0.indexed._C_ctype,
+        ffp0 = FieldFromPointer(f._C_field_dmap, f._C_symbol)
+        ffp1 = FieldFromPointer(f._C_field_arity, f._C_symbol)
+        args = [*call.arguments, ca.index, ffp1]
+        args[args.index(f.dmap)] = Cast(ffp0, dtype=f.c0.indexed._C_ctype,
                                         reinterpret=True)
         call1 = call._rebuild(arguments=args)
         efunc1 = Transformer({call: call1}).visit(efunc)
 
+        #TODO Propagate index through call stack until haloupdate...
+
         # Store the new Callables
         processed.update({kernel1.name: kernel1, k: efunc1})
 

devito/types/array.py

@@ -223,7 +223,7 @@ class MappedArrayMixin:
     _C_field_shape = 'shape'
     _C_field_size = 'size'
     _C_field_nbytes = 'nbytes'
-    _C_field_sizeofelem = 'sizeofelem'
+    _C_field_arity = 'arity'
 
     _C_size_type = c_uint64
 
@@ -233,7 +233,7 @@ class MappedArrayMixin:
                                           (_C_field_shape, POINTER(_C_size_type)),
                                           (_C_field_size, _C_size_type),
                                           (_C_field_nbytes, _C_size_type),
-                                          (_C_field_sizeofelem, _C_size_type)]}))
+                                          (_C_field_arity, _C_size_type)]}))
 
 
 class ArrayMapped(MappedArrayMixin, Array):