kernelcreation.py 5.86 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
from types import MappingProxyType
2
from pystencils.assignment_collection import AssignmentCollection
Martin Bauer's avatar
Martin Bauer committed
3
from pystencils.gpucuda.indexing import indexing_creator_from_params
Martin Bauer's avatar
Martin Bauer committed
4
5


Martin Bauer's avatar
Martin Bauer committed
6
7
def create_kernel(equations, target='cpu', data_type="double", iteration_slice=None, ghost_layers=None,
                  cpu_openmp=False, cpu_vectorize_info=None,
Martin Bauer's avatar
Martin Bauer committed
8
                  gpu_indexing='block', gpu_indexing_params=MappingProxyType({})):
Martin Bauer's avatar
Martin Bauer committed
9
10
    """
    Creates abstract syntax tree (AST) of kernel, using a list of update equations.
11
    :param equations: either be a plain list of equations or a AssignmentCollection object
Martin Bauer's avatar
Martin Bauer committed
12
    :param target: 'cpu', 'llvm' or 'gpu'
Martin Bauer's avatar
Martin Bauer committed
13
    :param data_type: data type used for all untyped symbols (i.e. non-fields), can also be a dict from symbol name
Martin Bauer's avatar
Martin Bauer committed
14
                     to type
Martin Bauer's avatar
Martin Bauer committed
15
16
    :param iteration_slice: rectangular subset to iterate over, if not specified the complete non-ghost layer \
                            part of the field is iterated over
Martin Bauer's avatar
Martin Bauer committed
17
    :param ghost_layers: if left to default, the number of necessary ghost layers is determined automatically
Martin Bauer's avatar
Martin Bauer committed
18
                        a single integer specifies the ghost layer count at all borders, can also be a sequence of
Martin Bauer's avatar
Martin Bauer committed
19
                        pairs [(x_lower_gl, x_upper_gl), .... ]
Martin Bauer's avatar
Martin Bauer committed
20
21

    CPU specific Parameters:
Martin Bauer's avatar
Martin Bauer committed
22
23
    :param cpu_openmp: True or number of threads for OpenMP parallelization, False for no OpenMP
    :param cpu_vectorize_info: pair of instruction set name ('sse, 'avx', 'avx512') and data type ('float', 'double')
Martin Bauer's avatar
Martin Bauer committed
24
25

    GPU specific Parameters
Martin Bauer's avatar
Martin Bauer committed
26
27
    :param gpu_indexing: either 'block' or 'line' , or custom indexing class (see gpucuda/indexing.py)
    :param gpu_indexing_params: dict with indexing parameters (constructor parameters of indexing class)
Martin Bauer's avatar
Martin Bauer committed
28
                              e.g. for 'block' one can specify {'block_size': (20, 20, 10) }
Martin Bauer's avatar
Martin Bauer committed
29
30
31
32
33
34

    :return: abstract syntax tree object, that can either be printed as source code or can be compiled with
             through its compile() function
    """

    # ----  Normalizing parameters
Martin Bauer's avatar
Martin Bauer committed
35
    split_groups = ()
36
    if isinstance(equations, AssignmentCollection):
Martin Bauer's avatar
Martin Bauer committed
37
38
        if 'split_groups' in equations.simplification_hints:
            split_groups = equations.simplification_hints['split_groups']
Martin Bauer's avatar
Martin Bauer committed
39
        equations = equations.all_assignments
Martin Bauer's avatar
Martin Bauer committed
40
41
42

    # ----  Creating ast
    if target == 'cpu':
Martin Bauer's avatar
Martin Bauer committed
43
44
45
46
47
48
49
        from pystencils.cpu import create_kernel
        from pystencils.cpu import add_openmp
        ast = create_kernel(equations, type_info=data_type, split_groups=split_groups,
                            iteration_slice=iteration_slice, ghost_layers=ghost_layers)
        if cpu_openmp:
            add_openmp(ast, num_threads=cpu_openmp)
        if cpu_vectorize_info:
Martin Bauer's avatar
Martin Bauer committed
50
51
            import pystencils.backends.simd_instruction_sets as vec
            from pystencils.vectorization import vectorize
Martin Bauer's avatar
Martin Bauer committed
52
            vec_params = cpu_vectorize_info
Martin Bauer's avatar
Martin Bauer committed
53
54
            vec.selected_instruction_set = vec.x86_vector_instruction_set(instruction_set=vec_params[0],
                                                                          data_type=vec_params[1])
Martin Bauer's avatar
Martin Bauer committed
55
56
57
            vectorize(ast)
        return ast
    elif target == 'llvm':
Martin Bauer's avatar
Martin Bauer committed
58
59
60
        from pystencils.llvm import create_kernel
        ast = create_kernel(equations, type_info=data_type, split_groups=split_groups,
                            iteration_slice=iteration_slice, ghost_layers=ghost_layers)
Martin Bauer's avatar
Martin Bauer committed
61
62
        return ast
    elif target == 'gpu':
Martin Bauer's avatar
Martin Bauer committed
63
64
65
66
        from pystencils.gpucuda import create_cuda_kernel
        ast = create_cuda_kernel(equations, type_info=data_type,
                                 indexing_creator=indexing_creator_from_params(gpu_indexing, gpu_indexing_params),
                                 iteration_slice=iteration_slice, ghost_layers=ghost_layers)
Martin Bauer's avatar
Martin Bauer committed
67
68
69
70
71
        return ast
    else:
        raise ValueError("Unknown target %s. Has to be one of 'cpu', 'gpu' or 'llvm' " % (target,))


Martin Bauer's avatar
Martin Bauer committed
72
def create_indexed_kernel(assignments, index_fields, target='cpu', data_type="double", coordinate_names=('x', 'y', 'z'),
Martin Bauer's avatar
Martin Bauer committed
73
                          cpu_openmp=True, gpu_indexing='block', gpu_indexing_params=MappingProxyType({})):
Martin Bauer's avatar
Martin Bauer committed
74
    """
Martin Bauer's avatar
Martin Bauer committed
75
    Similar to :func:`create_kernel`, but here not all cells of a field are updated but only cells with
Martin Bauer's avatar
Martin Bauer committed
76
77
    coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling.

Martin Bauer's avatar
Martin Bauer committed
78
    The coordinates are stored in a separated index_field, which is a one dimensional array with struct data type.
Martin Bauer's avatar
Martin Bauer committed
79
    This struct has to contain fields named 'x', 'y' and for 3D fields ('z'). These names are configurable with the
Martin Bauer's avatar
Martin Bauer committed
80
    'coordinate_names' parameter. The struct can have also other fields that can be read and written in the kernel, for
Martin Bauer's avatar
Martin Bauer committed
81
82
    example boundary parameters.

Martin Bauer's avatar
Martin Bauer committed
83
84
    index_fields: list of index fields, i.e. 1D fields with struct data type
    coordinate_names: name of the coordinate fields in the struct data type
Martin Bauer's avatar
Martin Bauer committed
85
86
    """

Martin Bauer's avatar
Martin Bauer committed
87
88
    if isinstance(assignments, AssignmentCollection):
        assignments = assignments.all_assignments
Martin Bauer's avatar
Martin Bauer committed
89
    if target == 'cpu':
Martin Bauer's avatar
Martin Bauer committed
90
91
92
93
94
95
        from pystencils.cpu import create_indexed_kernel
        from pystencils.cpu import add_openmp
        ast = create_indexed_kernel(assignments, index_fields=index_fields, type_info=data_type,
                                    coordinate_names=coordinate_names)
        if cpu_openmp:
            add_openmp(ast, num_threads=cpu_openmp)
Martin Bauer's avatar
Martin Bauer committed
96
97
98
99
        return ast
    elif target == 'llvm':
        raise NotImplementedError("Indexed kernels are not yet supported in LLVM backend")
    elif target == 'gpu':
Martin Bauer's avatar
Martin Bauer committed
100
        from pystencils.gpucuda import created_indexed_cuda_kernel
Martin Bauer's avatar
Martin Bauer committed
101
102
103
        idx_creator = indexing_creator_from_params(gpu_indexing, gpu_indexing_params)
        ast = created_indexed_cuda_kernel(assignments, index_fields, type_info=data_type,
                                          coordinate_names=coordinate_names, indexing_creator=idx_creator)
Martin Bauer's avatar
Martin Bauer committed
104
105
106
        return ast
    else:
        raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,))