diff --git a/pystencils/kernelcreation.py b/pystencils/kernelcreation.py index 7f50e38c2164eeea9ce226e9602feaeaddc25eea..c4d5273d4f0b303e600a8ead6378878e1fdf403f 100644 --- a/pystencils/kernelcreation.py +++ b/pystencils/kernelcreation.py @@ -1,3 +1,4 @@ +import functools import itertools from types import MappingProxyType @@ -10,7 +11,6 @@ from pystencils.gpucuda.indexing import indexing_creator_from_params from pystencils.simp.assignment_collection import AssignmentCollection from pystencils.transformations import ( loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel) -import functools def create_kernel(assignments, @@ -131,7 +131,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. @@ -181,7 +183,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, @@ -190,6 +192,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,))