codegen: New init_with_zero option for packs

Used to handle new Kernel requirement when the Kernel expects output
arguments to be zero on entry.

Fixes firedrakeproject/firedrake#1768.

Author: wence-

pyop2/codegen/builder.py

@@ -145,9 +145,10 @@ def emit_unpack_instruction(self, *, loop_indices=None):
 
 class GlobalPack(Pack):
 
-    def __init__(self, outer, access):
+    def __init__(self, outer, access, init_with_zero=False):
         self.outer = outer
         self.access = access
+        self.init_with_zero = init_with_zero
 
     def kernel_arg(self, loop_indices=None):
         pack = self.pack(loop_indices)
@@ -169,11 +170,15 @@ def pack(self, loop_indices=None):
         # vectorisation loop transformations privatise these reduction
         # variables. The extra memory movement cost is minimal.
         loop_indices = self.pick_loop_indices(*loop_indices)
-        if self.access in {INC, WRITE}:
+        if self.init_with_zero:
+            also_zero = {MIN, MAX}
+        else:
+            also_zero = set()
+        if self.access in {INC, WRITE} | also_zero:
             val = Zero((), self.outer.dtype)
             multiindex = MultiIndex(*(Index(e) for e in shape))
             self._pack = Materialise(PackInst(loop_indices), val, multiindex)
-        elif self.access in {READ, RW, MIN, MAX}:
+        elif self.access in {READ, RW, MIN, MAX} - also_zero:
             multiindex = MultiIndex(*(Index(e) for e in shape))
             expr = Indexed(self.outer, multiindex)
             self._pack = Materialise(PackInst(loop_indices), expr, multiindex)
@@ -203,13 +208,15 @@ def emit_unpack_instruction(self, *, loop_indices=None):
 
 class DatPack(Pack):
     def __init__(self, outer, access, map_=None, interior_horizontal=False,
-                 view_index=None, layer_bounds=None):
+                 view_index=None, layer_bounds=None,
+                 init_with_zero=False):
         self.outer = outer
         self.map_ = map_
         self.access = access
         self.interior_horizontal = interior_horizontal
         self.view_index = view_index
         self.layer_bounds = layer_bounds
+        self.init_with_zero = init_with_zero
 
     def _mask(self, map_):
         """Override this if the map_ needs a masking condition."""
@@ -245,11 +252,15 @@ def pack(self, loop_indices=None):
         if self.view_index is None:
             shape = shape + self.outer.shape[1:]
 
-        if self.access in {INC, WRITE}:
+        if self.init_with_zero:
+            also_zero = {MIN, MAX}
+        else:
+            also_zero = set()
+        if self.access in {INC, WRITE} | also_zero:
             val = Zero((), self.outer.dtype)
             multiindex = MultiIndex(*(Index(e) for e in shape))
             self._pack = Materialise(PackInst(), val, multiindex)
-        elif self.access in {READ, RW, MIN, MAX}:
+        elif self.access in {READ, RW, MIN, MAX} - also_zero:
             multiindex = MultiIndex(*(Index(e) for e in shape))
             expr, mask = self._rvalue(multiindex, loop_indices=loop_indices)
             if mask is not None:
@@ -577,8 +588,9 @@ def emit_unpack_instruction(self, *,
 
 class WrapperBuilder(object):
 
-    def __init__(self, *, iterset, iteration_region=None, single_cell=False,
+    def __init__(self, *, kernel, iterset, iteration_region=None, single_cell=False,
                  pass_layer_to_kernel=False, forward_arg_types=()):
+        self.kernel = kernel
         self.arguments = []
         self.argument_accesses = []
         self.packed_args = []
@@ -593,6 +605,10 @@ def __init__(self, *, iterset, iteration_region=None, single_cell=False,
         self.single_cell = single_cell
         self.forward_arguments = tuple(Argument((), fa, pfx="farg") for fa in forward_arg_types)
 
+    @property
+    def requires_zeroed_output_arguments(self):
+        return self.kernel.requires_zeroed_output_arguments
+
     @property
     def subset(self):
         return isinstance(self.iterset, Subset)
@@ -605,9 +621,6 @@ def extruded(self):
     def constant_layers(self):
         return self.extruded and self.iterset.constant_layers
 
-    def set_kernel(self, kernel):
-        self.kernel = kernel
-
     @cached_property
     def loop_extents(self):
         return (Argument((), IntType, name="start"),
@@ -722,7 +735,8 @@ def add_argument(self, arg):
                     shape = (None, *a.data.shape[1:])
                     argument = Argument(shape, a.data.dtype, pfx="mdat")
                     packs.append(a.data.pack(argument, arg.access, self.map_(a.map, unroll=a.unroll_map),
-                                             interior_horizontal=interior_horizontal))
+                                             interior_horizontal=interior_horizontal,
+                                             init_with_zero=self.requires_zeroed_output_arguments))
                     self.arguments.append(argument)
                 pack = MixedDatPack(packs, arg.access, arg.dtype, interior_horizontal=interior_horizontal)
                 self.packed_args.append(pack)
@@ -740,15 +754,17 @@ def add_argument(self, arg):
                                     pfx="dat")
                 pack = arg.data.pack(argument, arg.access, self.map_(arg.map, unroll=arg.unroll_map),
                                      interior_horizontal=interior_horizontal,
-                                     view_index=view_index)
+                                     view_index=view_index,
+                                     init_with_zero=self.requires_zeroed_output_arguments)
                 self.arguments.append(argument)
                 self.packed_args.append(pack)
                 self.argument_accesses.append(arg.access)
         elif arg._is_global:
             argument = Argument(arg.data.dim,
                                 arg.data.dtype,
                                 pfx="glob")
-            pack = GlobalPack(argument, arg.access)
+            pack = GlobalPack(argument, arg.access,
+                              init_with_zero=self.requires_zeroed_output_arguments)
             self.arguments.append(argument)
             self.packed_args.append(pack)
             self.argument_accesses.append(arg.access)

pyop2/sequential.py

@@ -148,10 +148,12 @@ def code_to_compile(self):
         from pyop2.codegen.builder import WrapperBuilder
         from pyop2.codegen.rep2loopy import generate
 
-        builder = WrapperBuilder(iterset=self._iterset, iteration_region=self._iteration_region, pass_layer_to_kernel=self._pass_layer_arg)
+        builder = WrapperBuilder(kernel=self._kernel,
+                                 iterset=self._iterset,
+                                 iteration_region=self._iteration_region,
+                                 pass_layer_to_kernel=self._pass_layer_arg)
         for arg in self._args:
             builder.add_argument(arg)
-        builder.set_kernel(self._kernel)
 
         wrapper = generate(builder)
         if self._iterset._extruded: