Merge pull request #643 from OP2/wence/fix/permuted-map

Wence/fix/permuted map

Author: wence-

pyop2/codegen/builder.py

@@ -30,32 +30,18 @@ class Map(object):
 
     __slots__ = ("values", "offset", "interior_horizontal",
                  "variable", "unroll", "layer_bounds",
-                 "prefetch", "permutation")
+                 "prefetch", "_pmap_count")
 
     def __init__(self, map_, interior_horizontal, layer_bounds,
-                 values=None, offset=None, unroll=False):
+                 offset=None, unroll=False):
         self.variable = map_.iterset._extruded and not map_.iterset.constant_layers
         self.unroll = unroll
         self.layer_bounds = layer_bounds
         self.interior_horizontal = interior_horizontal
         self.prefetch = {}
-        if values is not None:
-            raise RuntimeError
-            self.values = values
-            if map_.offset is not None:
-                assert offset is not None
-            self.offset = offset
-            return
-
         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)
@@ -64,6 +50,7 @@ def __init__(self, map_, interior_horizontal, layer_bounds,
 
         self.values = values
         self.offset = offset
+        self._pmap_count = itertools.count()
 
     @property
     def shape(self):
@@ -73,7 +60,7 @@ def shape(self):
     def dtype(self):
         return self.values.dtype
 
-    def indexed(self, multiindex, layer=None):
+    def indexed(self, multiindex, layer=None, permute=lambda x: x):
         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
@@ -82,10 +69,7 @@ def indexed(self, multiindex, layer=None):
             base_key = None
             if base_key not in self.prefetch:
                 j = Index()
-                if self.permutation is None:
-                    base = Indexed(self.values, (n, j))
-                else:
-                    base = Indexed(self.values, (n, Indexed(self.permutation, (j,))))
+                base = Indexed(self.values, (n, permute(j)))
                 self.prefetch[base_key] = Materialise(PackInst(), base, MultiIndex(j))
 
             base = self.prefetch[base_key]
@@ -112,26 +96,58 @@ 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
-            if self.permutation is None:
-                base = Indexed(self.values, (n, i))
-            else:
-                base = Indexed(self.values, (n, Indexed(self.permutation, (i,))))
+            base = Indexed(self.values, (n, permute(i)))
             return base, (f, i)
 
-    def indexed_vector(self, n, shape, layer=None):
+    def indexed_vector(self, n, shape, layer=None, permute=lambda x: x):
         shape = self.shape[1:] + shape
         if self.interior_horizontal:
             shape = (2, ) + shape
         else:
             shape = (1, ) + shape
         f, i, j = (Index(e) for e in shape)
-        base, (f, i) = self.indexed((n, i, f), layer=layer)
+        base, (f, i) = self.indexed((n, i, f), layer=layer, permute=permute)
         init = Sum(Product(base, Literal(numpy.int32(j.extent))), j)
         pack = Materialise(PackInst(), init, MultiIndex(f, i, j))
         multiindex = tuple(Index(e) for e in pack.shape)
         return Indexed(pack, multiindex), multiindex
 
 
+class PMap(Map):
+    __slots__ = ("permutation",)
+
+    def __init__(self, map_, permutation):
+        # Copy over properties
+        self.variable = map_.variable
+        self.unroll = map_.unroll
+        self.layer_bounds = map_.layer_bounds
+        self.interior_horizontal = map_.interior_horizontal
+        self.prefetch = {}
+        self.values = map_.values
+        self.offset = map_.offset
+        offset = map_.offset
+        # TODO: this is a hack, rep2loopy should be in charge of
+        # generating all names!
+        count = next(map_._pmap_count)
+        if offset is not None:
+            if offset.shape:
+                # Have a named literal
+                offset = offset.value[permutation]
+                offset = NamedLiteral(offset, parent=self.values, suffix=f"permutation{count}_offset")
+            else:
+                offset = map_.offset
+        self.offset = offset
+        self.permutation = NamedLiteral(permutation, parent=self.values, suffix=f"permutation{count}")
+
+    def indexed(self, multiindex, layer=None):
+        permute = lambda x: Indexed(self.permutation, (x,))
+        return super().indexed(multiindex, layer=layer, permute=permute)
+
+    def indexed_vector(self, n, shape, layer=None):
+        permute = lambda x: Indexed(self.permutation, (x,))
+        return super().indexed_vector(n, shape, layer=layer, permute=permute)
+
+
 class Pack(metaclass=ABCMeta):
 
     def pick_loop_indices(self, loop_index, layer_index=None, entity_index=None):
@@ -818,9 +834,13 @@ def map_(self, map_, unroll=False):
         try:
             return self.maps[key]
         except KeyError:
-            map_ = Map(map_, interior_horizontal,
-                       (self.bottom_layer, self.top_layer),
-                       unroll=unroll)
+            if isinstance(map_, PermutedMap):
+                imap = self.map_(map_.map_, unroll=unroll)
+                map_ = PMap(imap, map_.permutation)
+            else:
+                map_ = Map(map_, interior_horizontal,
+                           (self.bottom_layer, self.top_layer),
+                           unroll=unroll)
             self.maps[key] = map_
             return map_
 
@@ -854,7 +874,11 @@ def wrapper_args(self):
         args.extend(self.arguments)
         # maps are refcounted
         for map_ in self.maps.values():
-            args.append(map_.values)
+            # But we don't need to emit stuff for PMaps because they
+            # are a Map (already seen + a permutation [encoded in the
+            # indexing]).
+            if not isinstance(map_, PMap):
+                args.append(map_.values)
         return tuple(args)
 
     def kernel_call(self):

test/unit/test_indirect_loop.py

@@ -278,6 +278,45 @@ def test_mixed_non_mixed_dat_itspace(self, mdat, mmap, iterset):
         assert all(mdat[0].data == 1.0) and mdat[1].data == 4096.0
 
 
+def test_permuted_map():
+    fromset = op2.Set(1)
+    toset = op2.Set(4)
+    d1 = op2.Dat(op2.DataSet(toset, 1), dtype=np.int32)
+    d2 = op2.Dat(op2.DataSet(toset, 1), dtype=np.int32)
+    d1.data[:] = np.arange(4, dtype=np.int32)
+    k = op2.Kernel("""
+    void copy(int *to, const int * restrict from) {
+        for (int i = 0; i < 4; i++) { to[i] = from[i]; }
+    }""", "copy")
+    m1 = op2.Map(fromset, toset, 4, values=[1, 2, 3, 0])
+    m2 = op2.PermutedMap(m1, [3, 2, 0, 1])
+    op2.par_loop(k, fromset, d2(op2.WRITE, m2), d1(op2.READ, m1))
+    expect = np.empty_like(d1.data)
+    expect[m1.values[..., m2.permutation]] = d1.data[m1.values]
+    assert (d1.data == np.arange(4, dtype=np.int32)).all()
+    assert (d2.data == expect).all()
+
+
+def test_permuted_map_both():
+    fromset = op2.Set(1)
+    toset = op2.Set(4)
+    d1 = op2.Dat(op2.DataSet(toset, 1), dtype=np.int32)
+    d2 = op2.Dat(op2.DataSet(toset, 1), dtype=np.int32)
+    d1.data[:] = np.arange(4, dtype=np.int32)
+    k = op2.Kernel("""
+    void copy(int *to, const int * restrict from) {
+        for (int i = 0; i < 4; i++) { to[i] = from[i]; }
+    }""", "copy")
+    m1 = op2.Map(fromset, toset, 4, values=[0, 2, 1, 3])
+    m2 = op2.PermutedMap(m1, [3, 2, 1, 0])
+    m3 = op2.PermutedMap(m1, [0, 2, 3, 1])
+    op2.par_loop(k, fromset, d2(op2.WRITE, m2), d1(op2.READ, m3))
+    expect = np.empty_like(d1.data)
+    expect[m1.values[..., m2.permutation]] = d1.data[m1.values[..., m3.permutation]]
+    assert (d1.data == np.arange(4, dtype=np.int32)).all()
+    assert (d2.data == expect).all()
+
+
 if __name__ == '__main__':
     import os
     pytest.main(os.path.abspath(__file__))