kernelcreation.py 1.84 KB
Newer Older
Jan Hönig's avatar
Jan Hönig committed
1
from pystencils.transformations import insert_casts
Martin Bauer's avatar
Martin Bauer committed
2
from pystencils.llvm.llvmjit import make_python_function
3
4


Martin Bauer's avatar
Martin Bauer committed
5
6
def create_kernel(assignments, function_name="kernel", type_info=None, split_groups=(),
                  iteration_slice=None, ghost_layers=None):
7
8
9
10
11
    """
    Creates an abstract syntax tree for a kernel function, by taking a list of update rules.

    Loops are created according to the field accesses in the equations.

12
13
14
15
16
17
18
19
20
21
22
23
24
    Args:
    assignments: list of sympy equations, containing accesses to :class:`pystencils.field.Field`.
                 Defining the update rules of the kernel
    function_name: name of the generated function - only important if generated code is written out
                   type_info: a map from symbol name to a C type specifier. If not specified all symbols are assumed to
                   be of type 'double' except symbols which occur on the left hand side of equations where the
                   right hand side is a sympy Boolean which are assumed to be 'bool' .
    split_groups: Specification on how to split up inner loop into multiple loops. For details see
                  transformation :func:`pystencils.transformation.split_inner_loop`
    iteration_slice: if not None, iteration is done only over this slice of the field
    ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
                 if None, the number of ghost layers is determined automatically and assumed to be equal for a
                 all dimensions
25
26
27

    :return: :class:`pystencils.ast.KernelFunction` node
    """
Martin Bauer's avatar
Martin Bauer committed
28
29
    from pystencils.cpu import create_kernel
    code = create_kernel(assignments, function_name, type_info, split_groups, iteration_slice, ghost_layers)
30
31
32
    code.body = insert_casts(code.body)
    code._compile_function = make_python_function
    code._backend = 'llvm'
Jan Hoenig's avatar
Jan Hoenig committed
33
    return code