diff --git a/pyop2/codegen/builder.py b/pyop2/codegen/builder.py index 32bab1de7..414e1dac5 100644 --- a/pyop2/codegen/builder.py +++ b/pyop2/codegen/builder.py @@ -6,7 +6,7 @@ import numpy from loopy.types import OpaqueType from pyop2.global_kernel import (GlobalKernelArg, DatKernelArg, MixedDatKernelArg, - MatKernelArg, MixedMatKernelArg, PermutedMapKernelArg) + MatKernelArg, MixedMatKernelArg, PermutedMapKernelArg, ComposedMapKernelArg) from pyop2.codegen.representation import (Accumulate, Argument, Comparison, DummyInstruction, Extent, FixedIndex, FunctionCall, Index, Indexed, @@ -154,6 +154,28 @@ def indexed_vector(self, n, shape, layer=None): return super().indexed_vector(n, shape, layer=layer, permute=permute) +class CMap(Map): + + def __init__(self, *maps_): + # Copy over properties + self.variable = maps_[0].variable + self.unroll = maps_[0].unroll + self.layer_bounds = maps_[0].layer_bounds + self.interior_horizontal = maps_[0].interior_horizontal + self.prefetch = {} + self.values = maps_[0].values + self.offset = maps_[0].offset + self.maps_ = maps_ + + def indexed(self, multiindex, layer=None): + n, i, f = multiindex + n_ = n + for map_ in reversed(self.maps_): + if map_ is not self.maps_[0]: + n_, (_, _) = map_.indexed(MultiIndex(n_, FixedIndex(0), Index()), layer=None) + return self.maps_[0].indexed(MultiIndex(n_, i, f), layer=layer) + + class Pack(metaclass=ABCMeta): def pick_loop_indices(self, loop_index, layer_index=None, entity_index=None): @@ -835,6 +857,8 @@ def _add_map(self, map_, unroll=False): if isinstance(map_, PermutedMapKernelArg): imap = self._add_map(map_.base_map, unroll) map_ = PMap(imap, numpy.asarray(map_.permutation, dtype=IntType)) + elif isinstance(map_, ComposedMapKernelArg): + map_ = CMap(*(self._add_map(m, unroll) for m in map_.base_maps)) else: map_ = Map(interior_horizontal, (self.bottom_layer, self.top_layer), @@ -878,7 +902,8 @@ def wrapper_args(self): # 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): + # CMaps do not have their own arguments, either. + if not isinstance(map_, (PMap, CMap)): args.append(map_.values) return tuple(args) diff --git a/pyop2/global_kernel.py b/pyop2/global_kernel.py index 2e7339d21..6ef49dfa6 100644 --- a/pyop2/global_kernel.py +++ b/pyop2/global_kernel.py @@ -61,6 +61,26 @@ def cache_key(self): return type(self), self.base_map.cache_key, tuple(self.permutation) +@dataclass(eq=False, init=False) +class ComposedMapKernelArg: + """Class representing a composed map input to the kernel. + + :param base_maps: An arbitrary combination of :class:`MapKernelArg`s, :class:`PermutedMapKernelArg`s, and :class:`ComposedMapKernelArg`s. + """ + + def __init__(self, *base_maps): + self.base_maps = base_maps + + def __post_init__(self): + for m in self.base_maps: + if not isinstance(m, (MapKernelArg, PermutedMapKernelArg, ComposedMapKernelArg)): + raise TypeError("base_maps must be a combination of MapKernelArgs, PermutedMapKernelArgs, and ComposedMapKernelArgs") + + @property + def cache_key(self): + return type(self), tuple(m.cache_key for m in self.base_maps) + + @dataclass(frozen=True) class GlobalKernelArg: """Class representing a :class:`pyop2.types.Global` being passed to the kernel. diff --git a/pyop2/op2.py b/pyop2/op2.py index 98c09f727..1fe7f9d8a 100644 --- a/pyop2/op2.py +++ b/pyop2/op2.py @@ -41,7 +41,7 @@ from pyop2.types import ( Set, ExtrudedSet, MixedSet, Subset, DataSet, MixedDataSet, - Map, MixedMap, PermutedMap, Sparsity, Halo, + Map, MixedMap, PermutedMap, ComposedMap, Sparsity, Halo, Global, GlobalDataSet, Dat, MixedDat, DatView, Mat ) @@ -64,7 +64,7 @@ 'MixedSet', 'Subset', 'DataSet', 'GlobalDataSet', 'MixedDataSet', 'Halo', 'Dat', 'MixedDat', 'Mat', 'Global', 'Map', 'MixedMap', 'Sparsity', 'parloop', 'Parloop', 'ParLoop', 'par_loop', - 'DatView', 'PermutedMap'] + 'DatView', 'PermutedMap', 'ComposedMap'] _initialised = False diff --git a/pyop2/parloop.py b/pyop2/parloop.py index 8384268cf..2863ab88f 100644 --- a/pyop2/parloop.py +++ b/pyop2/parloop.py @@ -16,7 +16,7 @@ MatKernelArg, MixedMatKernelArg, GlobalKernel) from pyop2.local_kernel import LocalKernel, CStringLocalKernel, CoffeeLocalKernel, LoopyLocalKernel from pyop2.types import (Access, Global, Dat, DatView, MixedDat, Mat, Set, - MixedSet, ExtrudedSet, Subset, Map, MixedMap) + MixedSet, ExtrudedSet, Subset, Map, ComposedMap, MixedMap) from pyop2.utils import cached_property @@ -25,7 +25,10 @@ class ParloopArg(abc.ABC): @staticmethod def check_map(m): if configuration["type_check"]: - if m.iterset.total_size > 0 and len(m.values_with_halo) == 0: + if isinstance(m, ComposedMap): + for m_ in m.maps_: + ParloopArg.check_map(m_) + elif m.iterset.total_size > 0 and len(m.values_with_halo) == 0: raise MapValueError(f"{m} is not initialized") diff --git a/pyop2/sparsity.pyx b/pyop2/sparsity.pyx index a55ecaa62..c4d3f1cc9 100644 --- a/pyop2/sparsity.pyx +++ b/pyop2/sparsity.pyx @@ -51,7 +51,9 @@ cdef extern from "petsc.h": PETSC_INSERT_VALUES "INSERT_VALUES" int PetscCalloc1(size_t, void*) int PetscMalloc1(size_t, void*) + int PetscMalloc2(size_t, void*, size_t, void*) int PetscFree(void*) + int PetscFree2(void*,void*) int MatSetValuesBlockedLocal(PETSc.PetscMat, PetscInt, PetscInt*, PetscInt, PetscInt*, PetscScalar*, PetscInsertMode) int MatSetValuesLocal(PETSc.PetscMat, PetscInt, PetscInt*, PetscInt, PetscInt*, @@ -193,7 +195,9 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d PetscScalar zero = 0.0 PetscInt nrow, ncol PetscInt rarity, carity, tmp_rarity, tmp_carity - PetscInt[:, ::1] rmap, cmap + PetscInt[:, ::1] rmap, cmap, tempmap + PetscInt **rcomposedmaps = NULL, **ccomposedmaps = NULL + PetscInt nrcomposedmaps = 0, nccomposedmaps = 0, rset_entry, cset_entry PetscInt *rvals PetscInt *cvals PetscInt *roffset @@ -213,23 +217,52 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d set_size = pair[0].iterset.size if set_size == 0: continue - # Memoryviews require writeable buffers - rflag = set_writeable(pair[0]) - cflag = set_writeable(pair[1]) - # Map values - rmap = pair[0].values_with_halo - cmap = pair[1].values_with_halo + rflags = [] + cflags = [] + if isinstance(pair[0], op2.ComposedMap): + m = pair[0].flattened_maps[0] + rflags.append(set_writeable(m)) + rmap = m.values_with_halo + nrcomposedmaps = len(pair[0].flattened_maps) - 1 + else: + rflags.append(set_writeable(pair[0])) # Memoryviews require writeable buffers + rmap = pair[0].values_with_halo # Map values + if isinstance(pair[1], op2.ComposedMap): + m = pair[1].flattened_maps[0] + cflags.append(set_writeable(m)) + cmap = m.values_with_halo + nccomposedmaps = len(pair[1].flattened_maps) - 1 + else: + cflags.append(set_writeable(pair[1])) + cmap = pair[1].values_with_halo + # Handle ComposedMaps + CHKERR(PetscMalloc2(nrcomposedmaps, &rcomposedmaps, nccomposedmaps, &ccomposedmaps)) + for i in range(nrcomposedmaps): + m = pair[0].flattened_maps[1 + i] + rflags.append(set_writeable(m)) + tempmap = m.values_with_halo + rcomposedmaps[i] = &tempmap[0, 0] + for i in range(nccomposedmaps): + m = pair[1].flattened_maps[1 + i] + cflags.append(set_writeable(m)) + tempmap = m.values_with_halo + ccomposedmaps[i] = &tempmap[0, 0] # Arity of maps rarity = pair[0].arity carity = pair[1].arity - if not extruded: # The non-extruded case is easy, we just walk over the # rmap and cmap entries and set a block of values. CHKERR(PetscCalloc1(rarity*carity*rdim*cdim, &values)) for set_entry in range(set_size): - CHKERR(MatSetValuesBlockedLocal(mat.mat, rarity, &rmap[set_entry, 0], - carity, &cmap[set_entry, 0], + rset_entry = set_entry + cset_entry = set_entry + for i in range(nrcomposedmaps): + rset_entry = rcomposedmaps[nrcomposedmaps - 1 - i][rset_entry] + for i in range(nccomposedmaps): + cset_entry = ccomposedmaps[nccomposedmaps - 1 - i][cset_entry] + CHKERR(MatSetValuesBlockedLocal(mat.mat, rarity, &rmap[rset_entry, 0], + carity, &cmap[cset_entry, 0], values, PETSC_INSERT_VALUES)) else: # The extruded case needs a little more work. @@ -268,6 +301,12 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d for i in range(carity): coffset[i] = pair[1].offset[i] for set_entry in range(set_size): + rset_entry = set_entry + cset_entry = set_entry + for i in range(nrcomposedmaps): + rset_entry = rcomposedmaps[nrcomposedmaps - 1 - i][rset_entry] + for i in range(nccomposedmaps): + cset_entry = ccomposedmaps[nccomposedmaps - 1 - i][cset_entry] if constant_layers: layer_start = layers[0, 0] layer_end = layers[0, 1] - 1 @@ -287,15 +326,15 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d # In the case of tmp_rarity == rarity this is just: # - # rvals[i] = rmap[set_entry, i] + layer_start * roffset[i] + # rvals[i] = rmap[rset_entry, i] + layer_start * roffset[i] # # But this means less special casing. for i in range(tmp_rarity): - rvals[i] = rmap[set_entry, i % rarity] + \ + rvals[i] = rmap[rset_entry, i % rarity] + \ (layer_start - layer_bottom + i // rarity) * roffset[i % rarity] # Ditto for i in range(tmp_carity): - cvals[i] = cmap[set_entry, i % carity] + \ + cvals[i] = cmap[cset_entry, i % carity] + \ (layer_start - layer_bottom + i // carity) * coffset[i % carity] for layer in range(layer_start, layer_end): CHKERR(MatSetValuesBlockedLocal(mat.mat, tmp_rarity, rvals, @@ -310,6 +349,15 @@ def fill_with_zeros(PETSc.Mat mat not None, dims, maps, iteration_regions, set_d CHKERR(PetscFree(cvals)) CHKERR(PetscFree(roffset)) CHKERR(PetscFree(coffset)) - restore_writeable(pair[0], rflag) - restore_writeable(pair[1], cflag) + CHKERR(PetscFree2(rcomposedmaps, ccomposedmaps)) + if isinstance(pair[0], op2.ComposedMap): + for m, rflag in zip(pair[0].flattened_maps, rflags): + restore_writeable(m, rflag) + else: + restore_writeable(pair[0], rflags[0]) + if isinstance(pair[1], op2.ComposedMap): + for m, cflag in zip(pair[1].flattened_maps, cflags): + restore_writeable(m, cflag) + else: + restore_writeable(pair[1], cflags[0]) CHKERR(PetscFree(values)) diff --git a/pyop2/types/map.py b/pyop2/types/map.py index 5bb955380..7bb7536f4 100644 --- a/pyop2/types/map.py +++ b/pyop2/types/map.py @@ -149,6 +149,13 @@ def __le__(self, o): """self<=o if o equals self or self._parent <= o.""" return self == o + @utils.cached_property + def flattened_maps(self): + """Return all component maps. + + This is useful to flatten nested :class:`ComposedMap`s.""" + return (self, ) + class PermutedMap(Map): """Composition of a standard :class:`Map` with a constant permutation. @@ -173,6 +180,10 @@ class PermutedMap(Map): want two global-sized data structures. """ def __init__(self, map_, permutation): + if not isinstance(map_, Map): + raise TypeError("map_ must be a Map instance") + if isinstance(map_, ComposedMap): + raise NotImplementedError("PermutedMap of ComposedMap not implemented: simply permute before composing") self.map_ = map_ self.permutation = np.asarray(permutation, dtype=Map.dtype) assert (np.unique(permutation) == np.arange(map_.arity, dtype=Map.dtype)).all() @@ -192,6 +203,85 @@ def __getattr__(self, name): return getattr(self.map_, name) +class ComposedMap(Map): + """Composition of :class:`Map`s, :class:`PermutedMap`s, and/or :class:`ComposedMap`s. + + :arg maps_: The maps to compose. + + Where normally staging to element data is performed as + + .. code-block:: + + local[i] = global[map[i]] + + With a :class:`ComposedMap` we instead get + + .. code-block:: + + local[i] = global[maps_[0][maps_[1][maps_[2][...[i]]]]] + + This might be useful if the map you want can be represented by + a composition of existing maps. + """ + def __init__(self, *maps_, name=None): + if not all(isinstance(m, Map) for m in maps_): + raise TypeError("All maps must be Map instances") + for tomap, frommap in zip(maps_[:-1], maps_[1:]): + if tomap.iterset is not frommap.toset: + raise ex.MapTypeError("tomap.iterset must match frommap.toset") + if tomap.comm is not frommap.comm: + raise ex.MapTypeError("All maps needs to share a communicator") + if frommap.arity != 1: + raise ex.MapTypeError("frommap.arity must be 1") + self._iterset = maps_[-1].iterset + self._toset = maps_[0].toset + self.comm = self._toset.comm + self._arity = maps_[0].arity + # Don't call super().__init__() to avoid calling verify_reshape() + self._values = None + self.shape = (self._iterset.total_size, self._arity) + self._name = name or "cmap_#x%x" % id(self) + self._offset = maps_[0]._offset + # A cache for objects built on top of this map + self._cache = {} + self.maps_ = tuple(maps_) + + @utils.cached_property + def _kernel_args_(self): + return tuple(itertools.chain(*[m._kernel_args_ for m in self.maps_])) + + @utils.cached_property + def _wrapper_cache_key_(self): + return tuple(m._wrapper_cache_key_ for m in self.maps_) + + @utils.cached_property + def _global_kernel_arg(self): + from pyop2.global_kernel import ComposedMapKernelArg + + return ComposedMapKernelArg(*(m._global_kernel_arg for m in self.maps_)) + + @utils.cached_property + def values(self): + raise RuntimeError("ComposedMap does not store values directly") + + @utils.cached_property + def values_with_halo(self): + raise RuntimeError("ComposedMap does not store values directly") + + def __str__(self): + return "OP2 ComposedMap of Maps: [%s]" % ",".join([str(m) for m in self.maps_]) + + def __repr__(self): + return "ComposedMap(%s)" % ",".join([repr(m) for m in self.maps_]) + + def __le__(self, o): + raise NotImplementedError("__le__ not implemented for ComposedMap") + + @utils.cached_property + def flattened_maps(self): + return tuple(itertools.chain(*(m.flattened_maps for m in self.maps_))) + + class MixedMap(Map, caching.ObjectCached): r"""A container for a bag of :class:`Map`\s.""" @@ -315,3 +405,7 @@ def __str__(self): def __repr__(self): return "MixedMap(%r)" % (self._maps,) + + @utils.cached_property + def flattened_maps(self): + raise NotImplementedError("flattend_maps should not be necessary for MixedMap") diff --git a/pyop2/types/mat.py b/pyop2/types/mat.py index 723647edc..de89b1421 100644 --- a/pyop2/types/mat.py +++ b/pyop2/types/mat.py @@ -18,7 +18,7 @@ from pyop2.types.access import Access from pyop2.types.data_carrier import DataCarrier from pyop2.types.dataset import DataSet, GlobalDataSet, MixedDataSet -from pyop2.types.map import Map +from pyop2.types.map import Map, ComposedMap from pyop2.types.set import MixedSet, Set, Subset @@ -165,7 +165,7 @@ def _process_args(cls, dsets, maps, *, iteration_regions=None, name=None, nest=N if not isinstance(m, Map): raise ex.MapTypeError( "All maps must be of type map, not type %r" % type(m)) - if len(m.values_with_halo) == 0 and m.iterset.total_size > 0: + if not isinstance(m, ComposedMap) and len(m.values_with_halo) == 0 and m.iterset.total_size > 0: raise ex.MapValueError( "Unpopulated map values when trying to build sparsity.") # Make sure that the "to" Set of each map in a pair is the set of diff --git a/test/unit/test_indirect_loop.py b/test/unit/test_indirect_loop.py index 35921a3bd..ab77d182b 100644 --- a/test/unit/test_indirect_loop.py +++ b/test/unit/test_indirect_loop.py @@ -317,6 +317,162 @@ def test_permuted_map_both(): assert (d2.data == expect).all() +@pytest.mark.parametrize("permuted", ["none", "pre"]) +def test_composed_map_two_maps(permuted): + arity = 2 + setB = op2.Set(3) + nodesetB = op2.Set(6) + datB = op2.Dat(op2.DataSet(nodesetB, 1), dtype=np.float64) + mapB = op2.Map(setB, nodesetB, arity, values=[[0, 1], [2, 3], [4, 5]]) + setA = op2.Set(5) + nodesetA = op2.Set(8) + datA = op2.Dat(op2.DataSet(nodesetA, 1), dtype=np.float64) + datA.data[:] = np.array([.0, .1, .2, .3, .4, .5, .6, .7], dtype=np.float64) + mapA0 = op2.Map(setA, nodesetA, arity, values=[[0, 1], [2, 3], [4, 5], [6, 7], [0, 1]]) + if permuted == "pre": + mapA0 = op2.PermutedMap(mapA0, [1, 0]) + mapA1 = op2.Map(setB, setA, 1, values=[3, 1, 2]) + mapA = op2.ComposedMap(mapA0, mapA1) + # "post" permutation is currently not supported + k = op2.Kernel(""" + void copy(double *to, const double * restrict from) { + for (int i = 0; i < 2; ++i) { to[i] = from[i]; } + }""", "copy") + op2.par_loop(k, setB, datB(op2.WRITE, mapB), datA(op2.READ, mapA)) + if permuted == "none": + assert (datB.data == np.array([.6, .7, .2, .3, .4, .5], dtype=np.float64)).all() + else: + assert (datB.data == np.array([.7, .6, .3, .2, .5, .4], dtype=np.float64)).all() + + +@pytest.mark.parametrize("nested", ["none", "first", "last"]) +@pytest.mark.parametrize("subset", [False, True]) +def test_composed_map_three_maps(nested, subset): + arity = 2 + setC = op2.Set(2) + nodesetC = op2.Set(4) + datC = op2.Dat(op2.DataSet(nodesetC, 1), dtype=np.float64) + mapC = op2.Map(setC, nodesetC, arity, values=[[0, 1], [2, 3]]) + setB = op2.Set(3) + setA = op2.Set(5) + nodesetA = op2.Set(8) + datA = op2.Dat(op2.DataSet(nodesetA, 1), dtype=np.float64) + datA.data[:] = np.array([.0, .1, .2, .3, .4, .5, .6, .7], dtype=np.float64) + mapA0 = op2.Map(setA, nodesetA, arity, values=[[0, 1], [2, 3], [4, 5], [6, 7], [0, 1]]) + mapA1 = op2.Map(setB, setA, 1, values=[3, 1, 2]) + mapA2 = op2.Map(setC, setB, 1, values=[2, 0]) + if nested == "none": + mapA = op2.ComposedMap(mapA0, mapA1, mapA2) + elif nested == "first": + mapA = op2.ComposedMap(op2.ComposedMap(mapA0, mapA1), mapA2) + elif nested == "last": + mapA = op2.ComposedMap(mapA0, op2.ComposedMap(mapA1, mapA2)) + else: + raise ValueError(f"Unknown nested param: {nested}") + k = op2.Kernel(""" + void copy(double *to, const double * restrict from) { + for (int i = 0; i < 2; ++i) { to[i] = from[i]; } + }""", "copy") + if subset: + indices = np.array([1], dtype=np.int32) + setC = op2.Subset(setC, indices) + op2.par_loop(k, setC, datC(op2.WRITE, mapC), datA(op2.READ, mapA)) + if subset: + assert (datC.data == np.array([.0, .0, .6, .7], dtype=np.float64)).all() + else: + assert (datC.data == np.array([.4, .5, .6, .7], dtype=np.float64)).all() + + +@pytest.mark.parametrize("variable", [False, True]) +@pytest.mark.parametrize("subset", [False, True]) +def test_composed_map_extrusion(variable, subset): + # variable: False + # + # +14-+-9-+-4-+ + # |13 | 8 | 3 | + # +12-+-7-+-2-+ + # |11 | 6 | 1 | + # +10-+-5-+-0-+ + # + # 0 1 2 <- setA + # 0 1 <- setC + # + # variable: True + # + # +12-+-7-+-4-+ + # |11 | 6 | 3 | + # +10-+-5-+-2-+ + # | 9 | | 1 | + # +-8-+ +-0-+ + # + # 0 1 2 <- setA + # 0 1 <- setC + # + arity = 3 + if variable: + # A layer is a copy of base layer, so cell_layer_index + 1 + layersC = [[1, 2 + 1], [0, 2 + 1]] + setC = op2.ExtrudedSet(op2.Set(2), layersC) + nodesetC = op2.Set(8) + datC = op2.Dat(op2.DataSet(nodesetC, 1), dtype=np.float64) + mapC = op2.Map(setC, nodesetC, arity, + values=[[5, 6, 7], + [0, 1, 2]], + offset=[2, 2, 2]) + layersA = [[0, 2 + 1], [1, 2 + 1], [0, 2 + 1]] + setA = op2.ExtrudedSet(op2.Set(3), layersA) + nodesetA = op2.Set(13) + datA = op2.Dat(op2.DataSet(nodesetA, 1), dtype=np.float64) + datA.data[:] = np.arange(0, 13, dtype=np.float64) + mapA0 = op2.Map(setA, nodesetA, arity, + values=[[8, 9, 10], + [5, 6, 7], + [0, 1, 2]], + offset=[2, 2, 2]) + mapA1 = op2.Map(setC, setA, 1, values=[1, 2]) + mapA = op2.ComposedMap(mapA0, mapA1) + if subset: + expected = np.array([0., 1., 2., 3., 4., 0., 0., 0.], dtype=np.float64) + else: + expected = np.array([0., 1., 2., 3., 4., 5., 6., 7.], dtype=np.float64) + else: + # A layer is a copy of base layer, so cell_layer_index + 1 + layersC = 2 + 1 + setC = op2.ExtrudedSet(op2.Set(2), layersC) + nodesetC = op2.Set(10) + datC = op2.Dat(op2.DataSet(nodesetC, 1), dtype=np.float64) + mapC = op2.Map(setC, nodesetC, arity, + values=[[5, 6, 7], + [0, 1, 2]], + offset=[2, 2, 2]) + layersA = 2 + 1 + setA = op2.ExtrudedSet(op2.Set(3), layersA) + nodesetA = op2.Set(15) + datA = op2.Dat(op2.DataSet(nodesetA, 1), dtype=np.float64) + datA.data[:] = np.arange(0, 15, dtype=np.float64) + mapA0 = op2.Map(setA, nodesetA, arity, + values=[[10, 11, 12], + [5, 6, 7], + [0, 1, 2]], + offset=[2, 2, 2]) + mapA1 = op2.Map(setC, setA, 1, values=[1, 2]) + mapA = op2.ComposedMap(mapA0, mapA1) + if subset: + expected = np.array([0., 1., 2., 3., 4., 0., 0., 0., 0., 0.], dtype=np.float64) + else: + expected = np.array([0., 1., 2., 3., 4., 5., 6., 7., 8., 9.], dtype=np.float64) + k = op2.Kernel(""" + void copy(double *to, const double * restrict from) { + for (int i = 0; i < 3; ++i) { to[i] = from[i]; } + }""", "copy") + if subset: + indices = np.array([1], dtype=np.int32) + setC = op2.Subset(setC, indices) + op2.par_loop(k, setC, datC(op2.WRITE, mapC), datA(op2.READ, mapA)) + print(datC.data) + assert (datC.data == expected).all() + + if __name__ == '__main__': import os pytest.main(os.path.abspath(__file__))