diff --git a/pystencils/kernelcreation.py b/pystencils/kernelcreation.py index 16b0c64246e5540621210fe8502963be06db709f..a2dd55352548a4a28399920599394d53b5205465 100644 --- a/pystencils/kernelcreation.py +++ b/pystencils/kernelcreation.py @@ -1,4 +1,5 @@ -from itertools import combinations +import functools +import itertools from types import MappingProxyType import sympy as sp @@ -12,7 +13,6 @@ from pystencils.simp.assignment_collection import AssignmentCollection from pystencils.stencil import direction_string_to_offset, inverse_direction_string from pystencils.transformations import ( loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel) -import functools def create_kernel(assignments, @@ -143,7 +143,9 @@ def create_indexed_kernel(assignments, cpu_openmp=True, gpu_indexing='block', gpu_indexing_params=MappingProxyType({}), - use_textures_for_interpolation=True): + use_textures_for_interpolation=True, + opencl_queue=None, + opencl_ctx=None): """ Similar to :func:`create_kernel`, but here not all cells of a field are updated but only cells with coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling. @@ -193,7 +195,7 @@ def create_indexed_kernel(assignments, return ast elif target == 'llvm': raise NotImplementedError("Indexed kernels are not yet supported in LLVM backend") - elif target == 'gpu': + elif target == 'gpu' or target == 'opencl': from pystencils.gpucuda import created_indexed_cuda_kernel idx_creator = indexing_creator_from_params(gpu_indexing, gpu_indexing_params) ast = created_indexed_cuda_kernel(assignments, @@ -202,6 +204,12 @@ def create_indexed_kernel(assignments, coordinate_names=coordinate_names, indexing_creator=idx_creator, use_textures_for_interpolation=use_textures_for_interpolation) + if target == 'opencl': + from pystencils.opencl.opencljit import make_python_function + ast._backend = 'opencl' + ast.compile = functools.partial(make_python_function, ast, opencl_queue, opencl_ctx) + ast._target = 'opencl' + ast._backend = 'opencl' return ast else: raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,)) @@ -298,7 +306,7 @@ def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions= outer_assignment = None conditions = {direction: condition(direction) for direction in stencil} for num_conditions in range(len(stencil)): - for combination in combinations(conditions.values(), num_conditions): + for combination in itertools.combinations(conditions.values(), num_conditions): for assignment in assignments: direction = stencil[assignment.lhs.index[0]] if conditions[direction] in combination: