Merge remote-tracking branch 'origin/master' into Melina97/after_changes

pyop2/base.py

@@ -815,6 +815,13 @@ def indices(self):
         """Returns the indices pointing in the superset."""
         return self._indices
 
+    @cached_property
+    def owned_indices(self):
+        """Return the indices that correspond to the owned entities of the
+        superset.
+        """
+        return self.indices[self.indices < self.superset.size]
+
     @cached_property
     def layers_array(self):
         if self._superset.constant_layers:
@@ -1620,10 +1627,21 @@ def zero(self, subset=None):
         """Zero the data associated with this :class:`Dat`
 
         :arg subset: A :class:`Subset` of entries to zero (optional)."""
+<<<<<<< HEAD
         if subset is None:
             self.data[:] = 0
         else:
             self.data[subset.indices] = 0
+=======
+        # If there is no subset we can safely zero the halo values.
+        if subset is None:
+            self._data[:] = 0
+            self.halo_valid = True
+        elif subset.superset != self.dataset.set:
+            raise MapValueError("The subset and dataset are incompatible")
+        else:
+            self.data[subset.owned_indices] = 0
+>>>>>>> origin/master
 
     @collective
     def copy(self, other, subset=None):
@@ -1634,9 +1652,22 @@ def copy(self, other, subset=None):
         if other is self:
             return
         if subset is None:
+<<<<<<< HEAD
             other.data[:] = self.data_ro
         else:
             other.data[subset.indices] = self.data_ro[subset.indices]
+=======
+            # If the current halo is valid we can also copy these values across.
+            if self.halo_valid:
+                other._data[:] = self._data
+                other.halo_valid = True
+            else:
+                other.data[:] = self.data_ro
+        elif subset.superset != self.dataset.set:
+            raise MapValueError("The subset and dataset are incompatible")
+        else:
+            other.data[subset.owned_indices] = self.data_ro[subset.owned_indices]
+>>>>>>> origin/master
 
     def __iter__(self):
         """Yield self when iterated over."""
@@ -1848,13 +1879,17 @@ def __imul__(self, other):
 
     def __itruediv__(self, other):
         """Pointwise division or scaling of fields."""
+<<<<<<< HEAD
         from numbers import Number
         if isinstance(other, Number):
             self.data[:] /= other
         else:
             self._check_shape(other)
             np.true_divide(self.data[:], other.data_ro, out=self.data[:], casting="unsafe")
         return self
+=======
+        return self._iop(other, operator.itruediv)
+>>>>>>> origin/master
 
     @collective
     def global_to_local_begin(self, access_mode):
@@ -2638,6 +2673,41 @@ def __le__(self, o):
         return self == o
 
 
+class PermutedMap(Map):
+    """Composition of a standard :class:`Map` with a constant permutation.
+
+    :arg map_: The map to permute.
+    :arg permutation: The permutation of the map indices.
+
+    Where normally staging to element data is performed as
+
+    .. code-block::
+
+       local[i] = global[map[i]]
+
+    With a :class:`PermutedMap` we instead get
+
+    .. code-block::
+
+       local[i] = global[map[permutation[i]]]
+
+    This might be useful if your local kernel wants data in a
+    different order to the one that the map provides, and you don't
+    want two global-sized data structures.
+    """
+    def __init__(self, map_, permutation):
+        self.map_ = map_
+        self.permutation = np.asarray(permutation, dtype=Map.dtype)
+        assert (np.unique(permutation) == np.arange(map_.arity, dtype=Map.dtype)).all()
+
+    @cached_property
+    def _wrapper_cache_key_(self):
+        return super()._wrapper_cache_key_ + (tuple(self.permutation),)
+
+    def __getattr__(self, name):
+        return getattr(self.map_, name)
+
+
 class MixedMap(Map, ObjectCached):
     r"""A container for a bag of :class:`Map`\s."""
 

pyop2/codegen/builder.py

@@ -16,7 +16,7 @@
                                           When, Zero)
 from pyop2.datatypes import IntType
 from pyop2.op2 import (ALL, INC, MAX, MIN, ON_BOTTOM, ON_INTERIOR_FACETS,
-                       ON_TOP, READ, RW, WRITE, Subset)
+                       ON_TOP, READ, RW, WRITE, Subset, PermutedMap)
 from pyop2.utils import cached_property
 
 
@@ -30,7 +30,7 @@ class Map(object):
 
     __slots__ = ("values", "offset", "interior_horizontal",
                  "variable", "unroll", "layer_bounds",
-                 "prefetch")
+                 "prefetch", "permutation")
 
     def __init__(self, map_, interior_horizontal, layer_bounds,
                  values=None, offset=None, unroll=False):
@@ -50,11 +50,17 @@ def __init__(self, map_, interior_horizontal, layer_bounds,
         offset = map_.offset
         shape = (None, ) + map_.shape[1:]
         values = Argument(shape, dtype=map_.dtype, pfx="map")
+        if isinstance(map_, PermutedMap):
+            self.permutation = NamedLiteral(map_.permutation, parent=values, suffix="permutation")
+            if offset is not None:
+                offset = offset[map_.permutation]
+        else:
+            self.permutation = None
         if offset is not None:
             if len(set(map_.offset)) == 1:
                 offset = Literal(offset[0], casting=True)
             else:
-                offset = NamedLiteral(offset, name=values.name + "_offset")
+                offset = NamedLiteral(offset, parent=values, suffix="offset")
 
         self.values = values
         self.offset = offset
@@ -76,7 +82,10 @@ def indexed(self, multiindex, layer=None):
             base_key = None
             if base_key not in self.prefetch:
                 j = Index()
-                base = Indexed(self.values, (n, j))
+                if self.permutation is None:
+                    base = Indexed(self.values, (n, j))
+                else:
+                    base = Indexed(self.values, (n, Indexed(self.permutation, (j,))))
                 self.prefetch[base_key] = Materialise(PackInst(), base, MultiIndex(j))
 
             base = self.prefetch[base_key]
@@ -103,7 +112,10 @@ def indexed(self, multiindex, layer=None):
             return Indexed(self.prefetch[key], (f, i)), (f, i)
         else:
             assert f.extent == 1 or f.extent is None
-            base = Indexed(self.values, (n, i))
+            if self.permutation is None:
+                base = Indexed(self.values, (n, i))
+            else:
+                base = Indexed(self.values, (n, Indexed(self.permutation, (i,))))
             return base, (f, i)
 
     def indexed_vector(self, n, shape, layer=None):

pyop2/codegen/optimise.py

@@ -1,8 +1,7 @@
 from pyop2.codegen.node import traversal, reuse_if_untouched, Memoizer
 from functools import singledispatch
 from pyop2.codegen.representation import (Index, RuntimeIndex, Node,
-                                          FunctionCall, Variable, Argument,
-                                          NamedLiteral)
+                                          FunctionCall, Variable, Argument)
 
 
 def collect_indices(expressions):
@@ -90,7 +89,7 @@ def index_merger(instructions, cache=None):
 
 @singledispatch
 def _rename_node(node, self):
-    """Replace division with multiplication
+    """Rename nodes
 
     :param node: root of expression
     :param self: function for recursive calls
@@ -103,7 +102,7 @@ def _rename_node(node, self):
 
 @_rename_node.register(Index)
 def _rename_node_index(node, self):
-    name = self.replace.get(node, node.name)
+    name = self.renamer(node)
     return Index(extent=node.extent, name=name)
 
 
@@ -114,38 +113,25 @@ def _rename_node_func(node, self):
     return FunctionCall(node.name, node.label, node.access, free_indices, *children)
 
 
-@_rename_node.register(RuntimeIndex)
-def _rename_node_rtindex(node, self):
-    children = tuple(map(self, node.children))
-    name = self.replace.get(node, node.name)
-    return RuntimeIndex(*children, name=name)
-
-
-@_rename_node.register(NamedLiteral)
-def _rename_node_namedliteral(node, self):
-    name = self.replace.get(node, node.name)
-    return NamedLiteral(node.value, name)
-
-
 @_rename_node.register(Variable)
 def _rename_node_variable(node, self):
-    name = self.replace.get(node, node.name)
+    name = self.renamer(node)
     return Variable(name, node.shape, node.dtype)
 
 
 @_rename_node.register(Argument)
 def _rename_node_argument(node, self):
-    name = self.replace.get(node, node.name)
+    name = self.renamer(node)
     return Argument(node.shape, node.dtype, name=name)
 
 
-def rename_nodes(instructions, replace):
+def rename_nodes(instructions, renamer):
     """Rename the nodes in the instructions.
 
     :param instructions: Iterable of nodes.
-    :param replace: Dictionary matching old names to new names.
+    :param renamer: Function that maps nodes to new names
     :return: List of instructions with nodes renamed.
     """
     mapper = Memoizer(_rename_node)
-    mapper.replace = replace
+    mapper.renamer = renamer
     return list(map(mapper, instructions))

pyop2/codegen/rep2loopy.py

@@ -15,7 +15,6 @@
 from collections import OrderedDict, defaultdict
 from functools import singledispatch, reduce
 import itertools
-import re
 import operator
 
 from pyop2.codegen.node import traversal, Node, Memoizer, reuse_if_untouched
@@ -430,24 +429,35 @@ def generate(builder, wrapper_name=None):
     instructions = instructions + initialiser
     mapper.initialisers = [tuple(merger(i) for i in inits) for inits in mapper.initialisers]
 
+    def name_generator(prefix):
+        yield from (f"{prefix}{i}" for i in itertools.count())
+
     # rename indices and nodes (so that the counters start from zero)
-    pattern = re.compile(r"^([a-zA-Z_]+)([0-9]+)(_offset)?$")
-    replacements = {}
-    counter = defaultdict(itertools.count)
-    for node in traversal(instructions):
-        if isinstance(node, (Index, RuntimeIndex, Variable, Argument, NamedLiteral)):
-            match = pattern.match(node.name)
-            if match is None:
-                continue
-            prefix, _, postfix = match.groups()
-            if postfix is None:
-                postfix = ""
-            replacements[node] = "%s%d%s" % (prefix, next(counter[(prefix, postfix)]), postfix)
-
-    instructions = rename_nodes(instructions, replacements)
-    mapper.initialisers = [rename_nodes(inits, replacements) for inits in mapper.initialisers]
-    parameters.wrapper_arguments = rename_nodes(parameters.wrapper_arguments, replacements)
-    s, e = rename_nodes([mapper(e) for e in builder.layer_extents], replacements)
+    node_names = {}
+    node_namers = dict((cls, name_generator(prefix))
+                       for cls, prefix in [(Index, "i"), (Variable, "t")])
+
+    def renamer(expr):
+        if isinstance(expr, Argument):
+            if expr._name is not None:
+                # Some arguments have given names
+                return expr._name
+            else:
+                # Otherwise generate one with their given prefix.
+                namer = node_namers.setdefault((type(expr), expr.prefix),
+                                               name_generator(expr.prefix))
+        else:
+            namer = node_namers[type(expr)]
+        try:
+            return node_names[expr]
+        except KeyError:
+            return node_names.setdefault(expr, next(namer))
+
+    instructions = rename_nodes(instructions, renamer)
+    mapper.initialisers = [rename_nodes(inits, renamer)
+                           for inits in mapper.initialisers]
+    parameters.wrapper_arguments = rename_nodes(parameters.wrapper_arguments, renamer)
+    s, e = rename_nodes([mapper(e) for e in builder.layer_extents], renamer)
     parameters.layer_start = s.name
     parameters.layer_end = e.name