periodicity.py 2.22 KB
Newer Older
1
import numpy as np
2
from itertools import product
Martin Bauer's avatar
Martin Bauer committed
3

4
5
import pystencils.gpucuda
import pystencils.opencl
Martin Bauer's avatar
Martin Bauer committed
6
from pystencils import Assignment, Field
Martin Bauer's avatar
Martin Bauer committed
7
from pystencils.gpucuda.kernelcreation import create_cuda_kernel
Martin Bauer's avatar
Martin Bauer committed
8
from pystencils.slicing import get_periodic_boundary_src_dst_slices, normalize_slice
9
10


Martin Bauer's avatar
Martin Bauer committed
11
def create_copy_kernel(domain_size, from_slice, to_slice, index_dimensions=0, index_dim_shape=1, dtype=np.float64):
12
13
    """Copies a rectangular part of an array to another non-overlapping part"""

Martin Bauer's avatar
Martin Bauer committed
14
15
16
    f = Field.create_generic("pdfs", len(domain_size), index_dimensions=index_dimensions, dtype=dtype)
    normalized_from_slice = normalize_slice(from_slice, f.spatial_shape)
    normalized_to_slice = normalize_slice(to_slice, f.spatial_shape)
17

Martin Bauer's avatar
Martin Bauer committed
18
19
20
    offset = [s1.start - s2.start for s1, s2 in zip(normalized_from_slice, normalized_to_slice)]
    assert offset == [s1.stop - s2.stop for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \
        "Slices have to have same size"
21

Martin Bauer's avatar
Martin Bauer committed
22
    update_eqs = []
23
24
25
26
    if index_dimensions < 2:
        index_dim_shape = [index_dim_shape]
    for i in product(*[range(d) for d in index_dim_shape]):
        eq = Assignment(f(*i), f[tuple(offset)](*i))
Martin Bauer's avatar
Martin Bauer committed
27
        update_eqs.append(eq)
28

29
    ast = create_cuda_kernel(update_eqs, iteration_slice=to_slice, skip_independence_check=True)
30
    return ast
31
32


Martin Bauer's avatar
Martin Bauer committed
33
def get_periodic_boundary_functor(stencil, domain_size, index_dimensions=0, index_dim_shape=1, ghost_layers=1,
34
35
                                  thickness=None, dtype=float, target='gpu', opencl_queue=None, opencl_ctx=None):
    assert target in ['gpu', 'opencl']
Martin Bauer's avatar
Martin Bauer committed
36
    src_dst_slice_tuples = get_periodic_boundary_src_dst_slices(stencil, ghost_layers, thickness)
37
    kernels = []
38

Martin Bauer's avatar
Martin Bauer committed
39
    for src_slice, dst_slice in src_dst_slice_tuples:
40
41
42
43
        ast = create_copy_kernel(domain_size, src_slice, dst_slice, index_dimensions, index_dim_shape, dtype)
        if target == 'gpu':
            kernels.append(pystencils.gpucuda.make_python_function(ast))
        else:
44
45
            ast._target = 'opencl'
            ast._backend = 'opencl'
46
            kernels.append(pystencils.opencl.make_python_function(ast, opencl_queue, opencl_ctx))
47

Martin Bauer's avatar
Martin Bauer committed
48
    def functor(pdfs, **_):
49
50
51
52
        for kernel in kernels:
            kernel(pdfs=pdfs)

    return functor