Merge remote-tracking branch 'remotes/origin/wence/feature/codegen-updates' into vectorisation

Author: sv2518

pyop2/base.py

@@ -3287,6 +3287,8 @@ class Kernel(Cached):
         empty)
     :param ldargs: A list of arguments to pass to the linker when
         compiling this Kernel.
+    :param requires_zeroed_output_arguments: Does this kernel require the
+        output arguments to be zeroed on entry when called? (default no)
     :param cpp: Is the kernel actually C++ rather than C?  If yes,
         then compile with the C++ compiler (kernel is wrapped in
         extern C for linkage reasons).
@@ -3309,7 +3311,7 @@ class Kernel(Cached):
     @classmethod
     @validate_type(('name', str, NameTypeError))
     def _cache_key(cls, code, name, opts={}, include_dirs=[], headers=[],
-                   user_code="", ldargs=None, cpp=False):
+                   user_code="", ldargs=None, cpp=False, requires_zeroed_output_arguments=False):
         # Both code and name are relevant since there might be multiple kernels
         # extracting different functions from the same code
         # Also include the PyOP2 version, since the Kernel class might change
@@ -3323,15 +3325,15 @@ def _cache_key(cls, code, name, opts={}, include_dirs=[], headers=[],
             code.update_persistent_hash(key_hash, LoopyKeyBuilder())
             code = key_hash.hexdigest()
         hashee = (str(code) + name + str(sorted(opts.items())) + str(include_dirs)
-                  + str(headers) + version + str(ldargs) + str(cpp))
+                  + str(headers) + version + str(ldargs) + str(cpp) + str(requires_zeroed_output_arguments))
         return md5(hashee.encode()).hexdigest()
 
     @cached_property
     def _wrapper_cache_key_(self):
         return (self._key, )
 
     def __init__(self, code, name, opts={}, include_dirs=[], headers=[],
-                 user_code="", ldargs=None, cpp=False):
+                 user_code="", ldargs=None, cpp=False, requires_zeroed_output_arguments=False):
         # Protect against re-initialization when retrieved from cache
         if self._initialized:
             return
@@ -3346,6 +3348,7 @@ def __init__(self, code, name, opts={}, include_dirs=[], headers=[],
         assert isinstance(code, (str, Node, loopy.Program, loopy.LoopKernel))
         self._code = code
         self._initialized = True
+        self.requires_zeroed_output_arguments = requires_zeroed_output_arguments
 
     @property
     def name(self):

pyop2/codegen/builder.py

@@ -51,7 +51,10 @@ def __init__(self, map_, interior_horizontal, layer_bounds,
         shape = (None, ) + map_.shape[1:]
         values = Argument(shape, dtype=map_.dtype, pfx="map")
         if offset is not None:
-            offset = NamedLiteral(offset, name=values.name + "_offset")
+            if len(set(map_.offset)) == 1:
+                offset = Literal(offset[0], casting=True)
+            else:
+                offset = NamedLiteral(offset, name=values.name + "_offset")
 
         self.values = values
         self.offset = offset
@@ -68,21 +71,33 @@ def indexed(self, multiindex, layer=None):
         n, i, f = multiindex
         if layer is not None and self.offset is not None:
             # For extruded mesh, prefetch the indirections for each map, so that they don't
-            # need to be recomputed. Different f values need to be treated separately.
+            # need to be recomputed.
+            # First prefetch the base map (not dependent on layers)
+            base_key = None
+            if base_key not in self.prefetch:
+                j = Index()
+                base = Indexed(self.values, (n, j))
+                self.prefetch[base_key] = Materialise(PackInst(), base, MultiIndex(j))
+
+            base = self.prefetch[base_key]
+
+            # Now prefetch the extruded part of the map (inside the layer loop).
+            # This is necessary so loopy DTRT for MatSetValues
+            # Different f values need to be treated separately.
             key = f.extent
             if key is None:
                 key = 1
             if key not in self.prefetch:
                 bottom_layer, _ = self.layer_bounds
-                offset_extent, = self.offset.shape
-                j = Index(offset_extent)
-                base = Indexed(self.values, (n, j))
-                if f.extent:
-                    k = Index(f.extent)
-                else:
-                    k = Index(1)
+                k = Index(f.extent if f.extent is not None else 1)
                 offset = Sum(Sum(layer, Product(Literal(numpy.int32(-1)), bottom_layer)), k)
-                offset = Product(offset, Indexed(self.offset, (j,)))
+                j = Index()
+                # Inline map offsets where all entries are identical.
+                if self.offset.shape == ():
+                    offset = Product(offset, self.offset)
+                else:
+                    offset = Product(offset, Indexed(self.offset, (j,)))
+                base = Indexed(base, (j, ))
                 self.prefetch[key] = Materialise(PackInst(), Sum(base, offset), MultiIndex(k, j))
 
             return Indexed(self.prefetch[key], (f, i)), (f, i)
@@ -130,38 +145,78 @@ 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):
-        return Indexed(self.outer, (Index(e) for e in self.outer.shape))
+        pack = self.pack(loop_indices)
+        return Indexed(pack, (Index(e) for e in pack.shape))
 
     def emit_pack_instruction(self, *, loop_indices=None):
+        return ()
+
+    def pack(self, loop_indices=None):
+        if hasattr(self, "_pack"):
+            return self._pack
+
         shape = self.outer.shape
-        if self.access is WRITE:
-            zero = Zero((), self.outer.dtype)
+        if self.access is READ:
+            # No packing required
+            return self.outer
+        # We don't need to pack for memory layout, however packing
+        # globals that are written is required such that subsequent
+        # vectorisation loop transformations privatise these reduction
+        # variables. The extra memory movement cost is minimal.
+        loop_indices = self.pick_loop_indices(*loop_indices)
+        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} - also_zero:
             multiindex = MultiIndex(*(Index(e) for e in shape))
-            yield Accumulate(PackInst(), Indexed(self.outer, multiindex), zero)
+            expr = Indexed(self.outer, multiindex)
+            self._pack = Materialise(PackInst(loop_indices), expr, multiindex)
         else:
-            return ()
-
-    def pack(self, loop_indices=None):
-        return None
+            raise ValueError("Don't know how to initialise pack for '%s' access" % self.access)
+        return self._pack
 
     def emit_unpack_instruction(self, *, loop_indices=None):
-        return ()
+        pack = self.pack(loop_indices)
+        loop_indices = self.pick_loop_indices(*loop_indices)
+        if pack is None:
+            return ()
+        elif self.access is READ:
+            return ()
+        elif self.access in {INC, MIN, MAX}:
+            op = {INC: Sum,
+                  MIN: Min,
+                  MAX: Max}[self.access]
+            multiindex = tuple(Index(e) for e in pack.shape)
+            rvalue = Indexed(self.outer, multiindex)
+            yield Accumulate(UnpackInst(loop_indices), rvalue, op(rvalue, Indexed(pack, multiindex)))
+        else:
+            multiindex = tuple(Index(e) for e in pack.shape)
+            rvalue = Indexed(self.outer, multiindex)
+            yield Accumulate(UnpackInst(loop_indices), rvalue, Indexed(pack, multiindex))
 
 
 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."""
@@ -197,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:
@@ -529,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 = []
@@ -545,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)
@@ -557,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"),
@@ -674,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)
@@ -692,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/codegen/rep2loopy.py

@@ -231,7 +231,7 @@ def loop_nesting(instructions, deps, outer_inames, kernel_name):
             if isinstance(insn.children[1], (Zero, Literal)):
                 nesting[insn] = outer_inames
             else:
-                nesting[insn] = runtime_indices([insn])
+                nesting[insn] = runtime_indices([insn]) | runtime_indices(insn.label.within_inames)
         else:
             assert isinstance(insn, FunctionCall)
             if insn.name in (petsc_functions | {kernel_name}):
@@ -468,9 +468,9 @@ def generate(builder, wrapper_name=None):
     from coffee.base import Node
 
     if isinstance(kernel._code, loopy.LoopKernel):
+        from loopy.transform.callable import _match_caller_callee_argument_dimension_
         knl = kernel._code
         wrapper = loopy.register_callable_kernel(wrapper, knl)
-        from loopy.transform.callable import _match_caller_callee_argument_dimension_
         wrapper = _match_caller_callee_argument_dimension_(wrapper, knl.name)
     else:
         # kernel is a string, add it to preamble

pyop2/codegen/representation.py

@@ -9,7 +9,8 @@
 
 
 class InstructionLabel(object):
-    pass
+    def __init__(self, within_inames=()):
+        self.within_inames = tuple(w for w in within_inames if isinstance(w, Node))
 
 
 class PackInst(InstructionLabel):
@@ -99,6 +100,8 @@ def set_extent(self, value):
         elif self.extent != value:
             raise ValueError("Inconsistent index extents")
 
+    dtype = numpy.int32
+
 
 class FixedIndex(Terminal, Scalar):
     __slots__ = ("value", )
@@ -108,7 +111,9 @@ class FixedIndex(Terminal, Scalar):
 
     def __init__(self, value):
         assert isinstance(value, numbers.Integral)
-        self.value = int(value)
+        self.value = numpy.int32(value)
+
+    dtype = numpy.int32
 
 
 class RuntimeIndex(Scalar):
@@ -266,7 +271,7 @@ def __init__(self, a, b):
     @cached_property
     def dtype(self):
         a, b = self.children
-        return a.dtype
+        return numpy.find_common_type([], [a.dtype, b.dtype])
 
 
 class Sum(Scalar):
@@ -280,7 +285,7 @@ def __init__(self, a, b):
     @cached_property
     def dtype(self):
         a, b = self.children
-        return a.dtype
+        return numpy.find_common_type([], [a.dtype, b.dtype])
 
 
 class Product(Scalar):
@@ -294,7 +299,7 @@ def __init__(self, a, b):
     @cached_property
     def dtype(self):
         a, b = self.children
-        return a.dtype
+        return numpy.find_common_type([], [a.dtype, b.dtype])
 
 
 class Indexed(Scalar):
@@ -382,7 +387,7 @@ def __init__(self, name, shape, dtype):
 
 
 class DummyInstruction(Node):
-    __slots__ = ("children",)
+    __slots__ = ("children", "label")
     __front__ = ("label",)
 
     def __init__(self, label, *children):
@@ -391,17 +396,13 @@ def __init__(self, label, *children):
 
 
 class Accumulate(Node):
-    __slots__ = ("children",)
+    __slots__ = ("children", "label")
     __front__ = ("label",)
 
     def __init__(self, label, lvalue, rvalue):
         self.children = (lvalue, rvalue)
         self.label = label
 
-    def reconstruct(self, *args):
-        new = type(self)(*self._cons_args(args))
-        return new
-
 
 class FunctionCall(Node):
     __slots__ = ("name", "access", "free_indices", "label", "children")
@@ -417,7 +418,7 @@ def __init__(self, name, label, access, free_indices, *arguments):
 
 
 class Conditional(Scalar):
-    __slots__ = ("children")
+    __slots__ = ("children", )
 
     def __init__(self, condition, then, else_):
         assert not condition.shape

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: