compiler: Zero an Array whose out-of-DOMAIN entries are data

Author: mloubout

devito/passes/iet/definitions.py

@@ -22,8 +22,8 @@
 )
 from devito.tools import as_list, as_mapper, as_tuple, filter_sorted, flatten
 from devito.types import (
-    Array, ComponentAccess, CustomDimension, DeviceMap, DeviceRM, Dimension, Eq, Symbol,
-    size_t
+    Array, ComponentAccess, CustomDimension, DeviceMap, DeviceRM, Dimension, Eq, Pointer,
+    Symbol, size_t
 )
 
 __all__ = ['DataManager', 'DeviceAwareDataManager', 'Storage']
@@ -172,11 +172,14 @@ def _alloc_host_array_on_high_bw_mem(self, site, obj, storage, *args):
         memptr = VOID(Byref(obj._C_symbol), '**')
         alignment = obj._data_alignment
         nbytes = SizeOf(obj._C_typedata)*as_long(obj.size)
-        alloc = self.langbb['host-alloc'](memptr, alignment, nbytes)
+        allocs = [decl, self.langbb['host-alloc'](memptr, alignment, nbytes)]
+        if obj._is_zero_init:
+            allocs.append(self.langbb['host-memset'](obj._C_symbol, 0, nbytes))
+            storage.include(self.langbb['header-memcpy'])
 
         free = self.langbb['host-free'](obj._C_symbol)
 
-        storage.update(obj, site, allocs=(decl, alloc), frees=free)
+        storage.update(obj, site, allocs=tuple(allocs), frees=free)
 
     def _alloc_local_array_on_high_bw_mem(self, site, obj, storage, *args):
         """
@@ -579,11 +582,23 @@ def _alloc_local_array_on_high_bw_mem(self, site, obj, storage):
         dofree = self.langbb['device-free']
 
         nbytes = SizeOf(obj._C_typedata)*obj.size
-        init = doalloc(nbytes, deviceid, retobj=obj)
+        allocs = [doalloc(nbytes, deviceid, retobj=obj)]
+
+        if obj._is_zero_init:
+            # No language here has a device-side memset, so stage the zeros
+            # through a scratch host block; it's a one-off, at setup time
+            scratch = Pointer(name=self.sregistry.make_name(prefix='zeros'))
+            allocs.extend([
+                Call('calloc', (obj.size, SizeOf(obj._C_typedata)),
+                     retobj=scratch, cast=True),
+                self.langbb['memcpy-to-device'](obj._C_symbol, scratch, nbytes),
+                self.langbb['host-free'](scratch),
+            ])
+            storage.include(self.langbb['header-memcpy'])
 
         free = dofree(obj._C_name, deviceid)
 
-        storage.update(obj, site, allocs=init, frees=free)
+        storage.update(obj, site, allocs=tuple(allocs), frees=free)
 
     def _map_array_on_high_bw_mem(self, site, obj, storage):
         """
@@ -704,16 +719,19 @@ def process(self, graph):
 
 def make_zero_init(obj, rcompile, sregistry):
     cdims = []
-    for d, (h0, h1), s in zip(
-        obj.dimensions, obj._size_halo, obj.symbolic_shape, strict=True
+    for d, (h0, h1), (p0, p1), s in zip(
+        obj.dimensions, obj._size_halo, obj._size_padding, obj.symbolic_shape,
+        strict=True
     ):
         if d.is_NonlinearDerived:
-            assert h0 == h1 == 0
+            assert h0 == h1 == p0 == p1 == 0
             m = 0
             M = s - 1
         else:
-            m = d.symbolic_min - h0
-            M = d.symbolic_max + h1
+            # Spans the whole allocation, padding included, or the untouched
+            # entries would be left holding junk
+            m = d.symbolic_min - h0 - p0
+            M = d.symbolic_max + h1 + p1
         cdims.append(CustomDimension(name=d.name, parent=d,
                                      symbolic_min=m, symbolic_max=M))
 

devito/passes/iet/languages/C.py

@@ -56,6 +56,8 @@ class CBB(LangBB):
             Call('free', (i,)),
         'host-free-pin': lambda i:
             Call('free', (i,)),
+        'host-memset': lambda i, j, k:
+            Call('memset', (i, j, k)),
         'alloc-global-symbol': lambda i, j, k:
             Call('memcpy', (i, j, k))
     }

devito/types/basic.py

@@ -1351,6 +1351,12 @@ def _size_nodomain(self):
 
         return DimensionTuple(*sizes, getters=self.dimensions, left=left, right=right)
 
+    _is_zero_init = False
+    """
+    Whether the entries outside `self`'s DOMAIN carry meaningful data rather
+    than scratch, in which case the whole allocation must be zeroed upfront.
+    """
+
     @property
     def _is_reduction_ready(self):
         """

tests/test_operator.py

@@ -1394,6 +1394,26 @@ def test_conditional_declarations(self):
             assert i[0].is_Expression
             assert i[0].expr.rhs is init_value
 
+    def test_zero_init_array(self):
+        """
+        An Array whose entries outside the DOMAIN are data, rather than
+        scratch, is zeroed right after being allocated.
+        """
+        grid = Grid(shape=(4, 4))
+
+        class ZeroInitArray(Array):
+            _is_zero_init = True
+
+        a = ZeroInitArray(name='a', dimensions=grid.dimensions,
+                          dtype=grid.dtype, space='local')
+        b = Array(name='b', dimensions=grid.dimensions, dtype=grid.dtype,
+                  space='local')
+
+        f = Function(name='f', grid=grid)
+
+        assert 'memset(a' in str(Operator(Eq(f, a.indexify())))
+        assert 'memset(b' not in str(Operator(Eq(f, b.indexify())))
+
     def test_nested_scalar_assigns(self):
         grid = Grid(shape=(4, 4))