1import numpy as np 

2from itertools import product 

3 

4import pystencils.gpucuda 

5import pystencils.opencl 

6from pystencils import Assignment, Field 

7from pystencils.gpucuda.kernelcreation import create_cuda_kernel 

8from pystencils.slicing import get_periodic_boundary_src_dst_slices, normalize_slice 

9 

10 

11def create_copy_kernel(domain_size, from_slice, to_slice, index_dimensions=0, index_dim_shape=1, dtype=np.float64): 

12 """Copies a rectangular part of an array to another non-overlapping part""" 

13 

14 f = Field.create_generic("pdfs", len(domain_size), index_dimensions=index_dimensions, dtype=dtype) 

15 normalized_from_slice = normalize_slice(from_slice, f.spatial_shape) 

16 normalized_to_slice = normalize_slice(to_slice, f.spatial_shape) 

17 

18 offset = [s1.start - s2.start for s1, s2 in zip(normalized_from_slice, normalized_to_slice)] 

19 assert offset == [s1.stop - s2.stop for s1, s2 in zip(normalized_from_slice, normalized_to_slice)], \ 

20 "Slices have to have same size" 

21 

22 update_eqs = [] 

23 if index_dimensions < 2: 

24 index_dim_shape = [index_dim_shape] 

25 for i in product(*[range(d) for d in index_dim_shape]): 

26 eq = Assignment(f(*i), f[tuple(offset)](*i)) 

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 

33def get_periodic_boundary_functor(stencil, domain_size, index_dimensions=0, index_dim_shape=1, ghost_layers=1, 

34 thickness=None, dtype=float, target='gpu', opencl_queue=None, opencl_ctx=None): 

35 assert target in ['gpu', 'opencl'] 

36 src_dst_slice_tuples = get_periodic_boundary_src_dst_slices(stencil, ghost_layers, thickness) 

37 kernels = [] 

38 

39 for src_slice, dst_slice in src_dst_slice_tuples: 

40 ast = create_copy_kernel(domain_size, src_slice, dst_slice, index_dimensions, index_dim_shape, dtype) 

41 if target == 'gpu': 

42 kernels.append(pystencils.gpucuda.make_python_function(ast)) 

43 else: 

44 ast._target = 'opencl' 

45 ast._backend = 'opencl' 

46 kernels.append(pystencils.opencl.make_python_function(ast, opencl_queue, opencl_ctx)) 

47 

48 def functor(pdfs, **_): 

49 for kernel in kernels: 

50 kernel(pdfs=pdfs) 

51 

52 return functor