WIP WIP WIP [MASSIVE OVERHAUL IN PROGRESS, STILL BREAKING]

Author: FabioLuporini

devito/passes/iet/definitions.py

@@ -168,49 +168,55 @@ def _alloc_mapped_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
         decl = Definition(obj)
 
+        # Symbols/pointers
+        size = Symbol(name='size', dtype=np.uint64, is_const=True)
+        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)
+
         # Allocate the Array struct
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
         nbytes = SizeOf(obj._C_typedata)
-        allocs = [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 Array struct
+        init = [*[DummyExpr(IndexedPointer(ffp1, i), s)
+                  for i, s in enumerate(obj.symbolic_shape)],
+                DummyExpr(size, obj.size, init=True),
+                DummyExpr(ffp2, size),
+                DummyExpr(ffp3, size*SizeOf(obj.indexed._C_typedata))]
 
         # Allocate the underlying host data
-        ffp0 = FieldFromPointer(obj._C_field_data, obj._C_symbol)
         memptr = VOID(Byref(ffp0), '**')
-        allocs.append(self.langbb['host-alloc-pin'](memptr, alignment, nbytes_param))
-
-        # Initialize the Array 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)
+        alloc2 = self.langbb['host-alloc-pin'](memptr, alignment, ffp3)
 
+        # Free all of the allocated data
         frees = [self.langbb['host-free-pin'](ffp0),
+                 self.langbb['host-free'](ffp1),
                  self.langbb['host-free'](obj._C_symbol)]
 
         # Allocate the underlying device data, if required by the backend
-        alloc, free = self._make_dmap_allocfree(obj, nbytes_param)
-
-        # Chain together all allocs and frees
-        allocs = as_tuple(allocs) + as_tuple(alloc)
-        frees = as_tuple(free) + as_tuple(frees)
+        alloc_dmap, free_dmap = self._make_dmap_allocfree(obj, ffp3)
 
         ret = Return(obj._C_symbol)
 
         # Wrap everything in a Callable so that we can reuse the same code
         # for equivalent Array structs
         name = self.sregistry.make_name(prefix='alloc')
-        body = (decl, *allocs, init0, init1, ret)
+        body = (decl, alloc0, alloc1, *init, alloc2, *as_tuple(alloc_dmap), ret)
         efunc0 = make_callable(name, body, retval=obj)
-        args = list(efunc0.parameters)
-        args[args.index(nbytes_param)] = nbytes_arg
-        alloc = Call(name, args, retobj=obj)
+        alloc = Call(name, efunc0.parameters, retobj=obj)
 
         # Same story for the frees
         name = self.sregistry.make_name(prefix='free')
+        frees = as_tuple(free_dmap) + as_tuple(frees)
         efunc1 = make_callable(name, frees)
         free = Call(name, efunc1.parameters)
 

devito/types/array.py

@@ -1,4 +1,4 @@
-from ctypes import POINTER, Structure, c_void_p, c_ulong, c_uint64
+from ctypes import POINTER, Structure, c_void_p, c_uint64
 from functools import cached_property
 
 import numpy as np
@@ -202,19 +202,27 @@ def _make_pointer(self, dim):
         return PointerArray(name='p%s' % self.name, dimensions=dim, array=self)
 
 
-class ArrayMapped(Array):
+class MappedArrayMixin:
 
     _C_structname = 'array'
     _C_field_data = 'data'
     _C_field_dmap = 'dmap'
-    _C_field_nbytes = 'nbytes'
+    _C_field_shape = 'shape'
     _C_field_size = 'size'
+    _C_field_nbytes = 'nbytes'
+
+    _C_size_type = c_uint64
 
     _C_ctype = POINTER(type(_C_structname, (Structure,),
                             {'_fields_': [(_C_field_data, c_restrict_void_p),
-                                          (_C_field_nbytes, c_ulong),
                                           (_C_field_dmap, c_void_p),
-                                          (_C_field_size, c_uint64)]}))
+                                          (_C_field_shape, POINTER(_C_size_type)),
+                                          (_C_field_size, _C_size_type),
+                                          (_C_field_nbytes, _C_size_type)]}))
+
+
+class ArrayMapped(MappedArrayMixin, Array):
+    pass
 
 
 class ArrayObject(ArrayBasic):
@@ -343,7 +351,7 @@ def array(self):
         return self._array
 
 
-class Bundle(ArrayBasic):
+class Bundle(MappedArrayMixin, ArrayBasic):
 
     """
     Tensor symbol representing an unrolled vector of AbstractFunctions.
@@ -490,17 +498,12 @@ def __getitem__(self, index):
             raise ValueError("Expected %d or %d indices, got %d instead"
                              % (self.ndim, self.ndim + 1, len(index)))
 
-    _C_structname = ArrayMapped._C_structname
-    _C_field_data = ArrayMapped._C_field_data
-    _C_field_nbytes = ArrayMapped._C_field_nbytes
-    _C_field_dmap = ArrayMapped._C_field_dmap
-    _C_field_size = ArrayMapped._C_field_size
-
     @property
     def _C_ctype(self):
         if self._mem_mapped:
-            return ArrayMapped._C_ctype
+            return super()._C_ctype
         else:
+            #TODO DROP???
             return POINTER(dtype_to_ctype(self.dtype))