compiler: Add arity field to struct array

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -177,8 +177,12 @@ 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,
@@ -189,6 +193,7 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         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_arity, obj._C_symbol)
 
         # Allocate the Array struct
         memptr = VOID(Byref(obj._C_symbol), '**')
@@ -213,6 +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, arity_param))
 
         # Allocate the underlying host data
         memptr = VOID(Byref(ffp0), '**')
@@ -254,8 +260,12 @@ 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,
@@ -265,6 +275,7 @@ def _alloc_bundle_struct_on_high_bw_mem(self, site, obj, storage):
         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_arity, obj._C_symbol)
 
         # Allocate the Bundle struct
         memptr = VOID(Byref(obj._C_symbol), '**')
@@ -289,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),

devito/types/array.py

@@ -223,6 +223,7 @@ class MappedArrayMixin:
     _C_field_shape = 'shape'
     _C_field_size = 'size'
     _C_field_nbytes = 'nbytes'
+    _C_field_arity = 'arity'
 
     _C_size_type = c_uint64
 
@@ -231,7 +232,8 @@ class MappedArrayMixin:
                                           (_C_field_dmap, c_void_p),
                                           (_C_field_shape, POINTER(_C_size_type)),
                                           (_C_field_size, _C_size_type),
-                                          (_C_field_nbytes, _C_size_type)]}))
+                                          (_C_field_nbytes, _C_size_type),
+                                          (_C_field_arity, _C_size_type)]}))
 
 
 class ArrayMapped(MappedArrayMixin, Array):