kernelcreation.py 11.1 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
from types import MappingProxyType
2
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
3
4
from pystencils.assignment import Assignment
from pystencils.astnodes import LoopOverCoordinate, Conditional, Block, SympyAssignment
Martin Bauer's avatar
Martin Bauer committed
5
from pystencils.cpu.vectorization import vectorize
Martin Bauer's avatar
Martin Bauer committed
6
7
8
from pystencils.simp.assignment_collection import AssignmentCollection
from pystencils.gpucuda.indexing import indexing_creator_from_params
from pystencils.transformations import remove_conditionals_in_staggered_kernel
Martin Bauer's avatar
Martin Bauer committed
9
10


Martin Bauer's avatar
Martin Bauer committed
11
def create_kernel(assignments, target='cpu', data_type="double", iteration_slice=None, ghost_layers=None,
Martin Bauer's avatar
Martin Bauer committed
12
                  cpu_openmp=False, cpu_vectorize_info=None,
Martin Bauer's avatar
Martin Bauer committed
13
                  gpu_indexing='block', gpu_indexing_params=MappingProxyType({})):
Martin Bauer's avatar
Martin Bauer committed
14
15
    """
    Creates abstract syntax tree (AST) of kernel, using a list of update equations.
16
17

    Args:
Martin Bauer's avatar
Martin Bauer committed
18
        assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection`
19
20
21
22
23
24
25
        target: 'cpu', 'llvm' or 'gpu'
        data_type: data type used for all untyped symbols (i.e. non-fields), can also be a dict from symbol name
                  to type
        iteration_slice: rectangular subset to iterate over, if not specified the complete non-ghost layer \
                         part of the field is iterated over
        ghost_layers: if left to default, the number of necessary ghost layers is determined automatically
                     a single integer specifies the ghost layer count at all borders, can also be a sequence of
Martin Bauer's avatar
Martin Bauer committed
26
                     pairs ``[(x_lower_gl, x_upper_gl), .... ]``
27
28

        cpu_openmp: True or number of threads for OpenMP parallelization, False for no OpenMP
Martin Bauer's avatar
Martin Bauer committed
29
30
31
32
        cpu_vectorize_info: a dictionary with keys, 'vector_instruction_set', 'assume_aligned' and 'nontemporal'
                            for documentation of these parameters see vectorize function. Example:
                            '{'vector_instruction_set': 'avx512', 'assume_aligned': True, 'nontemporal':True}'
        gpu_indexing: either 'block' or 'line' , or custom indexing class, see `AbstractIndexing`
33
        gpu_indexing_params: dict with indexing parameters (constructor parameters of indexing class)
Martin Bauer's avatar
Martin Bauer committed
34
                             e.g. for 'block' one can specify '{'block_size': (20, 20, 10) }'
35
36

    Returns:
Martin Bauer's avatar
Martin Bauer committed
37
        abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
38
        can be compiled with through its 'compile()' member
Martin Bauer's avatar
Martin Bauer committed
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

    Example:
        >>> import pystencils as ps
        >>> import numpy as np
        >>> s, d = ps.fields('s, d: [2D]')
        >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
        >>> ast = ps.create_kernel(assignment, target='cpu', cpu_openmp=True)
        >>> kernel = ast.compile()
        >>> d_arr = np.zeros([5, 5])
        >>> kernel(d=d_arr, s=np.ones([5, 5]))
        >>> d_arr
        array([[0., 0., 0., 0., 0.],
               [0., 4., 4., 4., 0.],
               [0., 4., 4., 4., 0.],
               [0., 4., 4., 4., 0.],
               [0., 0., 0., 0., 0.]])
Martin Bauer's avatar
Martin Bauer committed
55
56
57
    """

    # ----  Normalizing parameters
Martin Bauer's avatar
Martin Bauer committed
58
    split_groups = ()
Martin Bauer's avatar
Martin Bauer committed
59
60
61
62
63
64
    if isinstance(assignments, AssignmentCollection):
        if 'split_groups' in assignments.simplification_hints:
            split_groups = assignments.simplification_hints['split_groups']
        assignments = assignments.all_assignments
    if isinstance(assignments, Assignment):
        assignments = [assignments]
Martin Bauer's avatar
Martin Bauer committed
65
66
67

    # ----  Creating ast
    if target == 'cpu':
Martin Bauer's avatar
Martin Bauer committed
68
69
        from pystencils.cpu import create_kernel
        from pystencils.cpu import add_openmp
Martin Bauer's avatar
Martin Bauer committed
70
        ast = create_kernel(assignments, type_info=data_type, split_groups=split_groups,
Martin Bauer's avatar
Martin Bauer committed
71
72
73
74
                            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
75
76
77
78
79
80
            if cpu_vectorize_info is True:
                vectorize(ast, vector_instruction_set='avx', assume_aligned=False, nontemporal=None)
            elif isinstance(cpu_vectorize_info, dict):
                vectorize(ast, **cpu_vectorize_info)
            else:
                raise ValueError("Invalid value for cpu_vectorize_info")
Martin Bauer's avatar
Martin Bauer committed
81
82
        return ast
    elif target == 'llvm':
Martin Bauer's avatar
Martin Bauer committed
83
        from pystencils.llvm import create_kernel
Martin Bauer's avatar
Martin Bauer committed
84
        ast = create_kernel(assignments, type_info=data_type, split_groups=split_groups,
Martin Bauer's avatar
Martin Bauer committed
85
                            iteration_slice=iteration_slice, ghost_layers=ghost_layers)
Martin Bauer's avatar
Martin Bauer committed
86
87
        return ast
    elif target == 'gpu':
Martin Bauer's avatar
Martin Bauer committed
88
        from pystencils.gpucuda import create_cuda_kernel
Martin Bauer's avatar
Martin Bauer committed
89
        ast = create_cuda_kernel(assignments, type_info=data_type,
Martin Bauer's avatar
Martin Bauer committed
90
91
                                 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
92
93
94
95
96
        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
97
def create_indexed_kernel(assignments, index_fields, target='cpu', data_type="double", coordinate_names=('x', 'y', 'z'),
Martin Bauer's avatar
Martin Bauer committed
98
                          cpu_openmp=True, gpu_indexing='block', gpu_indexing_params=MappingProxyType({})):
Martin Bauer's avatar
Martin Bauer committed
99
    """
Martin Bauer's avatar
Martin Bauer committed
100
    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
101
102
    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
103
    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
104
    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
105
    '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
106
107
    example boundary parameters.

Martin Bauer's avatar
Martin Bauer committed
108
109
    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
110

Martin Bauer's avatar
Martin Bauer committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
    Example:
        >>> import pystencils as ps
        >>> import numpy as np
        >>>
        >>> # Index field stores the indices of the cell to visit together with optional values
        >>> index_arr_dtype = np.dtype([('x', np.int32), ('y', np.int32), ('val', np.double)])
        >>> index_arr = np.array([(1, 1, 0.1), (2, 2, 0.2), (3, 3, 0.3)], dtype=index_arr_dtype)
        >>> idx_field = ps.fields(idx=index_arr)
        >>>
        >>> # Additional values  stored in index field can be accessed in the kernel as well
        >>> s, d = ps.fields('s, d: [2D]')
        >>> assignment = ps.Assignment(d[0,0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val'))
        >>> ast = create_indexed_kernel(assignment, [idx_field], coordinate_names=('x', 'y'))
        >>> kernel = ast.compile()
        >>> d_arr = np.zeros([5, 5])
        >>> kernel(s=np.ones([5, 5]), d=d_arr, idx=index_arr)
        >>> d_arr
        array([[0. , 0. , 0. , 0. , 0. ],
               [0. , 4.1, 0. , 0. , 0. ],
               [0. , 0. , 4.2, 0. , 0. ],
               [0. , 0. , 0. , 4.3, 0. ],
               [0. , 0. , 0. , 0. , 0. ]])
    """
    if isinstance(assignments, Assignment):
        assignments = [assignments]
    elif isinstance(assignments, AssignmentCollection):
Martin Bauer's avatar
Martin Bauer committed
137
        assignments = assignments.all_assignments
Martin Bauer's avatar
Martin Bauer committed
138
    if target == 'cpu':
Martin Bauer's avatar
Martin Bauer committed
139
140
141
142
143
144
        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
145
146
147
148
        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
149
        from pystencils.gpucuda import created_indexed_cuda_kernel
Martin Bauer's avatar
Martin Bauer committed
150
151
152
        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
153
154
155
        return ast
    else:
        raise ValueError("Unknown target %s. Has to be either 'cpu' or 'gpu'" % (target,))
156
157
158
159
160


def create_staggered_kernel(staggered_field, expressions, subexpressions=(), target='cpu', **kwargs):
    """Kernel that updates a staggered field.

Martin Bauer's avatar
Martin Bauer committed
161
162
    .. image:: /img/staggered_grid.svg

163
    Args:
164
165
166
        staggered_field: field where the first index coordinate defines the location of the staggered value
                can have 1 or 2 index coordinates, in case of of two index coordinates at every staggered location
                a vector is stored, expressions has to be a sequence of sequences then
Martin Bauer's avatar
Martin Bauer committed
167
168
                where e.g. ``f[0,0](0)`` is interpreted as value at the left cell boundary, ``f[1,0](0)`` the right cell
                boundary and ``f[0,0](1)`` the southern cell boundary etc.
169
        expressions: sequence of expressions of length dim, defining how the east, southern, (bottom) cell boundary
170
                     should be update.
171
172
173
        subexpressions: optional sequence of Assignments, that define subexpressions used in the main expressions
        target: 'cpu' or 'gpu'
        kwargs: passed directly to create_kernel, iteration slice and ghost_layers parameters are not allowed
174

175
    Returns:
176
        AST, see `create_kernel`
177
178
    """
    assert 'iteration_slice' not in kwargs and 'ghost_layers' not in kwargs
179
    assert staggered_field.index_dimensions in (1, 2), 'Staggered field must have one or two index dimensions'
180
181
182
183
184
    dim = staggered_field.spatial_dimensions

    counters = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)]
    conditions = [counters[i] < staggered_field.shape[i] - 1 for i in range(dim)]
    assert len(expressions) == dim
185
186
187
188
189
    if staggered_field.index_dimensions == 2:
        assert all(len(sublist) == len(expressions[0]) for sublist in expressions), \
            "If staggered field has two index dimensions expressions has to be a sequence of sequences of all the " \
            "same length."

190
191
192
    final_assignments = []
    for d in range(dim):
        cond = sp.And(*[conditions[i] for i in range(dim) if d != i])
193
194
195
196
197
198
199
200
        if staggered_field.index_dimensions == 1:
            assignments = [Assignment(staggered_field(d), expressions[d])]
            a_coll = AssignmentCollection(assignments, list(subexpressions)).new_filtered([staggered_field(d)])
        elif staggered_field.index_dimensions == 2:
            assert staggered_field.has_fixed_index_shape
            assignments = [Assignment(staggered_field(d, i), expr) for i, expr in enumerate(expressions[d])]
            a_coll = AssignmentCollection(assignments, list(subexpressions))
            a_coll = a_coll.new_filtered([staggered_field(d, i) for i in range(staggered_field.index_shape[1])])
201
202
        sp_assignments = [SympyAssignment(a.lhs, a.rhs) for a in a_coll.all_assignments]
        final_assignments.append(Conditional(cond, Block(sp_assignments)))
203

204
205
206
207
208
209
    ghost_layers = [(1, 0)] * dim

    ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target, **kwargs)
    if target == 'cpu':
        remove_conditionals_in_staggered_kernel(ast)
    return ast