diff --git a/pyop2/codegen/builder.py b/pyop2/codegen/builder.py index 50d57ca25..e09db09ab 100644 --- a/pyop2/codegen/builder.py +++ b/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): diff --git a/test/unit/test_indirect_loop.py b/test/unit/test_indirect_loop.py index b1f4e3cbe..35921a3bd 100644 --- a/test/unit/test_indirect_loop.py +++ b/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__))