transformations.py 44.8 KB
Newer Older
1
import warnings
2
from collections import defaultdict, OrderedDict, namedtuple
3
from copy import deepcopy
Martin Bauer's avatar
Martin Bauer committed
4
from types import MappingProxyType
5

6
7
import sympy as sp
from sympy.logic.boolalg import Boolean
8
from sympy.tensor import IndexedBase
9
from pystencils.assignment import Assignment
10
from pystencils.assignment_collection.nestedscopes import NestedScopes
Martin Bauer's avatar
Martin Bauer committed
11
from pystencils.field import Field, FieldType
Martin Bauer's avatar
Martin Bauer committed
12
from pystencils.data_types import TypedSymbol, PointerType, StructType, get_base_type, cast_func, \
13
    pointer_arithmetic_func, get_type_of_expression, collate_types, create_type
14
from pystencils.kernelparameters import FieldPointerSymbol
Martin Bauer's avatar
Martin Bauer committed
15
from pystencils.slicing import normalize_slice
Martin Bauer's avatar
Martin Bauer committed
16
import pystencils.astnodes as ast
17
18


Martin Bauer's avatar
Martin Bauer committed
19
def filtered_tree_iteration(node, node_type, stop_type=None):
20
    for arg in node.args:
Martin Bauer's avatar
Martin Bauer committed
21
        if isinstance(arg, node_type):
22
            yield arg
Martin Bauer's avatar
Martin Bauer committed
23
24
25
        elif stop_type and isinstance(node, stop_type):
            continue

Martin Bauer's avatar
Martin Bauer committed
26
        yield from filtered_tree_iteration(arg, node_type)
27
28


29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def unify_shape_symbols(body, common_shape, fields):
    """Replaces symbols for array sizes to ensure they are represented by the same unique symbol.

    When creating a kernel with variable array sizes, all passed arrays must have the same size.
    This is ensured when the kernel is called. Inside the kernel this means that only on symbol has to be used instead
    of one for each field. For example shape_arr1[0]  and shape_arr2[0] must be equal, so they should also be
    represented by the same symbol.

    Args:
        body: ast node, for the kernel part where substitutions is made, is modified in-place
        common_shape: shape of the field that was chosen
        fields: all fields whose shapes should be replaced by common_shape
    """
    substitutions = {}
    for field in fields:
        assert len(field.spatial_shape) == len(common_shape)
        if not field.has_fixed_shape:
            for common_shape_component, shape_component in zip(common_shape, field.spatial_shape):
                if shape_component != common_shape_component:
                    substitutions[shape_component] = common_shape_component
    if substitutions:
        body.subs(substitutions)


Martin Bauer's avatar
Martin Bauer committed
53
def get_common_shape(field_set):
54
55
    """Takes a set of pystencils Fields and returns their common spatial shape if it exists. Otherwise
    ValueError is raised"""
Martin Bauer's avatar
Martin Bauer committed
56
57
58
59
60
61
62
63
    nr_of_fixed_shaped_fields = 0
    for f in field_set:
        if f.has_fixed_shape:
            nr_of_fixed_shaped_fields += 1

    if nr_of_fixed_shaped_fields > 0 and nr_of_fixed_shaped_fields != len(field_set):
        fixed_field_names = ",".join([f.name for f in field_set if f.has_fixed_shape])
        var_field_names = ",".join([f.name for f in field_set if not f.has_fixed_shape])
64
        msg = "Mixing fixed-shaped and variable-shape fields in a single kernel is not possible\n"
Martin Bauer's avatar
Martin Bauer committed
65
        msg += "Variable shaped: %s \nFixed shaped:    %s" % (var_field_names, fixed_field_names)
66
67
        raise ValueError(msg)

Martin Bauer's avatar
Martin Bauer committed
68
69
70
71
    shape_set = set([f.spatial_shape for f in field_set])
    if nr_of_fixed_shaped_fields == len(field_set):
        if len(shape_set) != 1:
            raise ValueError("Differently sized field accesses in loop body: " + str(shape_set))
72

Martin Bauer's avatar
Martin Bauer committed
73
    shape = list(sorted(shape_set, key=lambda e: str(e[0])))[0]
74
75
76
    return shape


Martin Bauer's avatar
Martin Bauer committed
77
78
79
80
def make_loop_over_domain(body, function_name, iteration_slice=None, ghost_layers=None, loop_order=None):
    """Uses :class:`pystencils.field.Field.Access` to create (multiple) loops around given AST.

    Args:
81
        body: Block object with inner loop contents
Martin Bauer's avatar
Martin Bauer committed
82
83
84
85
86
87
88
89
90
        function_name: name of generated C function
        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
        loop_order: loop ordering from outer to inner loop (optimal ordering is same as layout)

    Returns:
        :class:`LoopOverCoordinate` instance with nested loops, ordered according to field layouts
91
92
    """
    # find correct ordering by inspecting participating FieldAccesses
Martin Bauer's avatar
Martin Bauer committed
93
    field_accesses = body.atoms(Field.Access)
94
95
    field_accesses = {e for e in field_accesses if not e.is_absolute_access}

Martin Bauer's avatar
Martin Bauer committed
96
97
98
99
100
101
102
    # exclude accesses to buffers from field_list, because buffers are treated separately
    field_list = [e.field for e in field_accesses if not FieldType.is_buffer(e.field)]
    fields = set(field_list)

    if loop_order is None:
        loop_order = get_optimal_loop_ordering(fields)

103
104
    shape = get_common_shape(fields)
    unify_shape_symbols(body, common_shape=shape, fields=fields)
Martin Bauer's avatar
Martin Bauer committed
105
106
107
108
109
110
111
112
113
114
115

    if iteration_slice is not None:
        iteration_slice = normalize_slice(iteration_slice, shape)

    if ghost_layers is None:
        required_ghost_layers = max([fa.required_ghost_layers for fa in field_accesses])
        ghost_layers = [(required_ghost_layers, required_ghost_layers)] * len(loop_order)
    if isinstance(ghost_layers, int):
        ghost_layers = [(ghost_layers, ghost_layers)] * len(loop_order)

    current_body = body
Martin Bauer's avatar
Martin Bauer committed
116
    for i, loop_coordinate in enumerate(reversed(loop_order)):
Martin Bauer's avatar
Martin Bauer committed
117
        if iteration_slice is None:
Martin Bauer's avatar
Martin Bauer committed
118
119
120
            begin = ghost_layers[loop_coordinate][0]
            end = shape[loop_coordinate] - ghost_layers[loop_coordinate][1]
            new_loop = ast.LoopOverCoordinate(current_body, loop_coordinate, begin, end, 1)
Martin Bauer's avatar
Martin Bauer committed
121
            current_body = ast.Block([new_loop])
Martin Bauer's avatar
Martin Bauer committed
122
        else:
Martin Bauer's avatar
Martin Bauer committed
123
            slice_component = iteration_slice[loop_coordinate]
Martin Bauer's avatar
Martin Bauer committed
124
125
            if type(slice_component) is slice:
                sc = slice_component
Martin Bauer's avatar
Martin Bauer committed
126
                new_loop = ast.LoopOverCoordinate(current_body, loop_coordinate, sc.start, sc.stop, sc.step)
Martin Bauer's avatar
Martin Bauer committed
127
                current_body = ast.Block([new_loop])
Martin Bauer's avatar
Martin Bauer committed
128
            else:
Martin Bauer's avatar
Martin Bauer committed
129
                assignment = ast.SympyAssignment(ast.LoopOverCoordinate.get_loop_counter_symbol(loop_coordinate),
Martin Bauer's avatar
Martin Bauer committed
130
131
                                                 sp.sympify(slice_component))
                current_body.insert_front(assignment)
132

Martin Bauer's avatar
Martin Bauer committed
133
    ast_node = ast.KernelFunction(current_body, ghost_layers=ghost_layers, function_name=function_name, backend='cpu')
Martin Bauer's avatar
Martin Bauer committed
134
    return ast_node
135
136


Martin Bauer's avatar
Martin Bauer committed
137
def create_intermediate_base_pointer(field_access, coordinates, previous_ptr):
Martin Bauer's avatar
Martin Bauer committed
138
    r"""
139
    Addressing elements in structured arrays is done with :math:`ptr\left[ \sum_i c_i \cdot s_i \right]`
Martin Bauer's avatar
Martin Bauer committed
140
141
142
143
144
    where :math:`c_i` is the coordinate value and :math:`s_i` the stride of a coordinate.
    The sum can be split up into multiple parts, such that parts of it can be pulled before loops.
    This function creates such an access for coordinates :math:`i \in \mbox{coordinates}`.
    Returns a new typed symbol, where the name encodes which coordinates have been resolved.

145
146
147
148
149
150
151
152
153
    Args:
        field_access: instance of :class:`pystencils.field.Field.Access` which provides strides and offsets
        coordinates: mapping of coordinate ids to its value, where stride*value is calculated
        previous_ptr: the pointer which is de-referenced

    Returns
        tuple with the new pointer symbol and the calculated offset

    Examples:
Martin Bauer's avatar
Martin Bauer committed
154
        >>> field = Field.create_generic('myfield', spatial_dimensions=2, index_dimensions=1)
Martin Bauer's avatar
Martin Bauer committed
155
        >>> x, y = sp.symbols("x y")
Martin Bauer's avatar
Martin Bauer committed
156
157
        >>> prev_pointer = TypedSymbol("ptr", "double")
        >>> create_intermediate_base_pointer(field[1,-2](5), {0: x}, prev_pointer)
158
        (ptr_01, _stride_myfield_0*x + _stride_myfield_0)
Martin Bauer's avatar
Martin Bauer committed
159
        >>> create_intermediate_base_pointer(field[1,-2](5), {0: x, 1 : y }, prev_pointer)
160
        (ptr_01_1m2, _stride_myfield_0*x + _stride_myfield_0 + _stride_myfield_1*y - 2*_stride_myfield_1)
Martin Bauer's avatar
Martin Bauer committed
161
    """
Martin Bauer's avatar
Martin Bauer committed
162
    field = field_access.field
163
164
    offset = 0
    name = ""
Martin Bauer's avatar
Martin Bauer committed
165
    list_to_hash = []
Martin Bauer's avatar
Martin Bauer committed
166
167
    for coordinate_id, coordinate_value in coordinates.items():
        offset += field.strides[coordinate_id] * coordinate_value
168

Martin Bauer's avatar
Martin Bauer committed
169
170
171
        if coordinate_id < field.spatial_dimensions:
            offset += field.strides[coordinate_id] * field_access.offsets[coordinate_id]
            if type(field_access.offsets[coordinate_id]) is int:
Martin Bauer's avatar
Martin Bauer committed
172
                name += "_%d%d" % (coordinate_id, field_access.offsets[coordinate_id])
173
            else:
Martin Bauer's avatar
Martin Bauer committed
174
                list_to_hash.append(field_access.offsets[coordinate_id])
175
        else:
Martin Bauer's avatar
Martin Bauer committed
176
            if type(coordinate_value) is int:
Martin Bauer's avatar
Martin Bauer committed
177
                name += "_%d%d" % (coordinate_id, coordinate_value)
178
            else:
Martin Bauer's avatar
Martin Bauer committed
179
                list_to_hash.append(coordinate_value)
180

Martin Bauer's avatar
Martin Bauer committed
181
    if len(list_to_hash) > 0:
Martin Bauer's avatar
Martin Bauer committed
182
        name += "_%0.6X" % (hash(tuple(list_to_hash)))
183

Martin Bauer's avatar
Martin Bauer committed
184
    name = name.replace("-", 'm')
Martin Bauer's avatar
Martin Bauer committed
185
186
    new_ptr = TypedSymbol(previous_ptr.name + name, previous_ptr.dtype)
    return new_ptr, offset
187
188


Martin Bauer's avatar
Martin Bauer committed
189
def parse_base_pointer_info(base_pointer_specification, loop_order, spatial_dimensions, index_dimensions):
190
    """
Martin Bauer's avatar
Martin Bauer committed
191
    Creates base pointer specification for :func:`resolve_field_accesses` function.
Martin Bauer's avatar
Martin Bauer committed
192
193
194

    Specification of how many and which intermediate pointers are created for a field access.
    For example [ (0), (2,3,)]  creates on base pointer for coordinates 2 and 3 and writes the offset for coordinate
195
    zero directly in the field access. These specifications are defined dependent on the loop ordering.
Martin Bauer's avatar
Martin Bauer committed
196
197
    This function translates more readable version into the specification above.

198
    Allowed specifications:
Martin Bauer's avatar
Martin Bauer committed
199
200
201
202
203
204
        - "spatialInner<int>" spatialInner0 is the innermost loop coordinate,
          spatialInner1 the loop enclosing the innermost
        - "spatialOuter<int>" spatialOuter0 is the outermost loop
        - "index<int>": index coordinate
        - "<int>": specifying directly the coordinate

Martin Bauer's avatar
Martin Bauer committed
205
206
207
    Args:
        base_pointer_specification: nested list with above specifications
        loop_order: list with ordering of loops from outer to inner
Martin Bauer's avatar
Martin Bauer committed
208
209
        spatial_dimensions: number of spatial dimensions
        index_dimensions: number of index dimensions
Martin Bauer's avatar
Martin Bauer committed
210
211
212

    Returns:
        list of tuples that can be passed to :func:`resolve_field_accesses`
Martin Bauer's avatar
Martin Bauer committed
213
214
215
216
217

    Examples:
        >>> parse_base_pointer_info([['spatialOuter0'], ['index0']], loop_order=[2,1,0],
        ...                         spatial_dimensions=3, index_dimensions=1)
        [[0], [3], [1, 2]]
218
219
    """
    result = []
Martin Bauer's avatar
Martin Bauer committed
220
221
    specified_coordinates = set()
    loop_order = list(reversed(loop_order))
Martin Bauer's avatar
Martin Bauer committed
222
    for spec_group in base_pointer_specification:
Martin Bauer's avatar
Martin Bauer committed
223
224
225
        new_group = []

        def add_new_element(elem):
Martin Bauer's avatar
Martin Bauer committed
226
            if elem >= spatial_dimensions + index_dimensions:
Martin Bauer's avatar
Martin Bauer committed
227
228
229
230
231
                raise ValueError("Coordinate %d does not exist" % (elem,))
            new_group.append(elem)
            if elem in specified_coordinates:
                raise ValueError("Coordinate %d specified two times" % (elem,))
            specified_coordinates.add(elem)
Martin Bauer's avatar
Martin Bauer committed
232
        for element in spec_group:
233
            if type(element) is int:
Martin Bauer's avatar
Martin Bauer committed
234
                add_new_element(element)
235
236
237
238
            elif element.startswith("spatial"):
                element = element[len("spatial"):]
                if element.startswith("Inner"):
                    index = int(element[len("Inner"):])
Martin Bauer's avatar
Martin Bauer committed
239
                    add_new_element(loop_order[index])
240
241
                elif element.startswith("Outer"):
                    index = int(element[len("Outer"):])
Martin Bauer's avatar
Martin Bauer committed
242
                    add_new_element(loop_order[-index])
243
                elif element == "all":
Martin Bauer's avatar
Martin Bauer committed
244
                    for i in range(spatial_dimensions):
Martin Bauer's avatar
Martin Bauer committed
245
                        add_new_element(i)
246
247
248
249
                else:
                    raise ValueError("Could not parse " + element)
            elif element.startswith("index"):
                index = int(element[len("index"):])
Martin Bauer's avatar
Martin Bauer committed
250
                add_new_element(spatial_dimensions + index)
251
252
253
            else:
                raise ValueError("Unknown specification %s" % (element,))

Martin Bauer's avatar
Martin Bauer committed
254
        result.append(new_group)
255

Martin Bauer's avatar
Martin Bauer committed
256
    all_coordinates = set(range(spatial_dimensions + index_dimensions))
Martin Bauer's avatar
Martin Bauer committed
257
    rest = all_coordinates - specified_coordinates
258
259
    if rest:
        result.append(list(rest))
260

261
262
263
    return result


Martin Bauer's avatar
Martin Bauer committed
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
def get_base_buffer_index(ast_node, loop_counters=None, loop_iterations=None):
    """Used for buffer fields to determine the linearized index of the buffer dependent on loop counter symbols.

    Args:
        ast_node: ast before any field accesses are resolved
        loop_counters: for CPU kernels: leave to default 'None' (can be determined from loop nodes)
                       for GPU kernels: list of 'loop counters' from inner to outer loop
        loop_iterations: number of iterations of each loop from inner to outer, for CPU kernels leave to default

    Returns:
        base buffer index - required by 'resolve_buffer_accesses' function
    """
    if loop_counters is None or loop_iterations is None:
        loops = [l for l in filtered_tree_iteration(ast_node, ast.LoopOverCoordinate, ast.SympyAssignment)]
        loops.reverse()
        parents_of_innermost_loop = list(parents_of_type(loops[0], ast.LoopOverCoordinate, include_current=True))
        assert len(loops) == len(parents_of_innermost_loop)
        assert all(l1 is l2 for l1, l2 in zip(loops, parents_of_innermost_loop))

        loop_iterations = [(l.stop - l.start) / l.step for l in loops]
        loop_counters = [l.loop_counter_symbol for l in loops]

    field_accesses = ast_node.atoms(Field.Access)
    buffer_accesses = {fa for fa in field_accesses if FieldType.is_buffer(fa.field)}
    loop_counters = [v * len(buffer_accesses) for v in loop_counters]

    base_buffer_index = loop_counters[0]
    stride = 1
    for idx, var in enumerate(loop_counters[1:]):
        cur_stride = loop_iterations[idx]
        stride *= int(cur_stride) if isinstance(cur_stride, float) else cur_stride
        base_buffer_index += var * stride
    return base_buffer_index


Martin Bauer's avatar
Martin Bauer committed
299
def resolve_buffer_accesses(ast_node, base_buffer_index, read_only_field_names=set()):
Martin Bauer's avatar
Martin Bauer committed
300

Martin Bauer's avatar
Martin Bauer committed
301
    def visit_sympy_expr(expr, enclosing_block, sympy_assignment):
302
        if isinstance(expr, Field.Access):
Martin Bauer's avatar
Martin Bauer committed
303
            field_access = expr
304
305

            # Do not apply transformation if field is not a buffer
Martin Bauer's avatar
Martin Bauer committed
306
            if not FieldType.is_buffer(field_access.field):
307
308
                return expr

Martin Bauer's avatar
Martin Bauer committed
309
            buffer = field_access.field
310
            field_ptr = FieldPointerSymbol(buffer.name, buffer.dtype, const=buffer.name in read_only_field_names)
311

Martin Bauer's avatar
Martin Bauer committed
312
313
            buffer_index = base_buffer_index
            if len(field_access.index) > 1:
314
315
                raise RuntimeError('Only indexing dimensions up to 1 are currently supported in buffers!')

Martin Bauer's avatar
Martin Bauer committed
316
317
318
            if len(field_access.index) > 0:
                cell_index = field_access.index[0]
                buffer_index += cell_index
319

Martin Bauer's avatar
Martin Bauer committed
320
321
            result = ast.ResolvedFieldAccess(field_ptr, buffer_index, field_access.field, field_access.offsets,
                                             field_access.index)
322

Martin Bauer's avatar
Martin Bauer committed
323
            return visit_sympy_expr(result, enclosing_block, sympy_assignment)
324
325
326
327
        else:
            if isinstance(expr, ast.ResolvedFieldAccess):
                return expr

Martin Bauer's avatar
Martin Bauer committed
328
            new_args = [visit_sympy_expr(e, enclosing_block, sympy_assignment) for e in expr.args]
329
            kwargs = {'evaluate': False} if type(expr) in (sp.Add, sp.Mul, sp.Piecewise) else {}
Martin Bauer's avatar
Martin Bauer committed
330
331
332
333
334
335
336
337
            return expr.func(*new_args, **kwargs) if new_args else expr

    def visit_node(sub_ast):
        if isinstance(sub_ast, ast.SympyAssignment):
            enclosing_block = sub_ast.parent
            assert type(enclosing_block) is ast.Block
            sub_ast.lhs = visit_sympy_expr(sub_ast.lhs, enclosing_block, sub_ast)
            sub_ast.rhs = visit_sympy_expr(sub_ast.rhs, enclosing_block, sub_ast)
338
        else:
Martin Bauer's avatar
Martin Bauer committed
339
340
            for i, a in enumerate(sub_ast.args):
                visit_node(a)
341

Martin Bauer's avatar
Martin Bauer committed
342
    return visit_node(ast_node)
343

344

Martin Bauer's avatar
Martin Bauer committed
345
def resolve_field_accesses(ast_node, read_only_field_names=set(),
Martin Bauer's avatar
Martin Bauer committed
346
347
                           field_to_base_pointer_info=MappingProxyType({}),
                           field_to_fixed_coordinates=MappingProxyType({})):
Martin Bauer's avatar
Martin Bauer committed
348
349
350
    """
    Substitutes :class:`pystencils.field.Field.Access` nodes by array indexing

351
352
353
354
355
356
    Args:
        ast_node: the AST root
        read_only_field_names: set of field names which are considered read-only
        field_to_base_pointer_info: a list of tuples indicating which intermediate base pointers should be created
                                    for details see :func:`parse_base_pointer_info`
        field_to_fixed_coordinates: map of field name to a tuple of coordinate symbols. Instead of using the loop
Martin Bauer's avatar
Martin Bauer committed
357
                                    counters to index the field these symbols are used as coordinates
358
359
360

    Returns
        transformed AST
Martin Bauer's avatar
Martin Bauer committed
361
    """
Martin Bauer's avatar
Martin Bauer committed
362
363
    field_to_base_pointer_info = OrderedDict(sorted(field_to_base_pointer_info.items(), key=lambda pair: pair[0]))
    field_to_fixed_coordinates = OrderedDict(sorted(field_to_fixed_coordinates.items(), key=lambda pair: pair[0]))
364

Martin Bauer's avatar
Martin Bauer committed
365
    def visit_sympy_expr(expr, enclosing_block, sympy_assignment):
366
        if isinstance(expr, Field.Access):
Martin Bauer's avatar
Martin Bauer committed
367
            field_access = expr
Martin Bauer's avatar
Martin Bauer committed
368
            field = field_access.field
369

Martin Bauer's avatar
Martin Bauer committed
370
            if field_access.indirect_addressing_fields:
371
372
                new_offsets = tuple(visit_sympy_expr(off, enclosing_block, sympy_assignment)
                                    for off in field_access.offsets)
Martin Bauer's avatar
Martin Bauer committed
373
374
375
376
377
                new_indices = tuple(visit_sympy_expr(ind, enclosing_block, sympy_assignment)
                                    if isinstance(ind, sp.Basic) else ind
                                    for ind in field_access.index)
                field_access = Field.Access(field_access.field, new_offsets,
                                            new_indices, field_access.is_absolute_access)
378

Martin Bauer's avatar
Martin Bauer committed
379
380
            if field.name in field_to_base_pointer_info:
                base_pointer_info = field_to_base_pointer_info[field.name]
381
            else:
Martin Bauer's avatar
Martin Bauer committed
382
                base_pointer_info = [list(range(field.index_dimensions + field.spatial_dimensions))]
383

384
            field_ptr = FieldPointerSymbol(field.name, field.dtype, const=field.name in read_only_field_names)
385

Martin Bauer's avatar
Martin Bauer committed
386
387
388
389
            def create_coordinate_dict(group_param):
                coordinates = {}
                for e in group_param:
                    if e < field.spatial_dimensions:
Martin Bauer's avatar
Martin Bauer committed
390
                        if field.name in field_to_fixed_coordinates:
391
392
393
394
                            if not field_access.is_absolute_access:
                                coordinates[e] = field_to_fixed_coordinates[field.name][e]
                            else:
                                coordinates[e] = 0
395
                        else:
Martin Bauer's avatar
Martin Bauer committed
396
397
398
399
                            if not field_access.is_absolute_access:
                                coordinates[e] = ast.LoopOverCoordinate.get_loop_counter_symbol(e)
                            else:
                                coordinates[e] = 0
Martin Bauer's avatar
Martin Bauer committed
400
                        coordinates[e] *= field.dtype.item_size
401
                    else:
402
                        if isinstance(field.dtype, StructType):
Martin Bauer's avatar
Martin Bauer committed
403
                            assert field.index_dimensions == 1
Martin Bauer's avatar
Martin Bauer committed
404
405
                            accessed_field_name = field_access.index[0]
                            assert isinstance(accessed_field_name, str)
Martin Bauer's avatar
Martin Bauer committed
406
                            coordinates[e] = field.dtype.get_element_offset(accessed_field_name)
407
                        else:
Martin Bauer's avatar
Martin Bauer committed
408
                            coordinates[e] = field_access.index[e - field.spatial_dimensions]
409

Martin Bauer's avatar
Martin Bauer committed
410
                return coordinates
411

Martin Bauer's avatar
Martin Bauer committed
412
            last_pointer = field_ptr
413

Martin Bauer's avatar
Martin Bauer committed
414
415
            for group in reversed(base_pointer_info[1:]):
                coord_dict = create_coordinate_dict(group)
Martin Bauer's avatar
Martin Bauer committed
416
                new_ptr, offset = create_intermediate_base_pointer(field_access, coord_dict, last_pointer)
Martin Bauer's avatar
Martin Bauer committed
417
418
419
420
                if new_ptr not in enclosing_block.symbols_defined:
                    new_assignment = ast.SympyAssignment(new_ptr, last_pointer + offset, is_const=False)
                    enclosing_block.insert_before(new_assignment, sympy_assignment)
                last_pointer = new_ptr
421

Martin Bauer's avatar
Martin Bauer committed
422
            coord_dict = create_coordinate_dict(base_pointer_info[0])
Martin Bauer's avatar
Martin Bauer committed
423
            _, offset = create_intermediate_base_pointer(field_access, coord_dict, last_pointer)
Martin Bauer's avatar
Martin Bauer committed
424
425
            result = ast.ResolvedFieldAccess(last_pointer, offset, field_access.field,
                                             field_access.offsets, field_access.index)
426

Martin Bauer's avatar
Martin Bauer committed
427
428
            if isinstance(get_base_type(field_access.field.dtype), StructType):
                new_type = field_access.field.dtype.get_element_type(field_access.index[0])
Martin Bauer's avatar
Martin Bauer committed
429
                result = cast_func(result, new_type)
430

Martin Bauer's avatar
Martin Bauer committed
431
            return visit_sympy_expr(result, enclosing_block, sympy_assignment)
432
        else:
Martin Bauer's avatar
Martin Bauer committed
433
434
            if isinstance(expr, ast.ResolvedFieldAccess):
                return expr
435

Martin Bauer's avatar
Martin Bauer committed
436
            new_args = [visit_sympy_expr(e, enclosing_block, sympy_assignment) for e in expr.args]
437
            kwargs = {'evaluate': False} if type(expr) in (sp.Add, sp.Mul, sp.Piecewise) else {}
Martin Bauer's avatar
Martin Bauer committed
438
439
440
441
442
443
444
445
            return expr.func(*new_args, **kwargs) if new_args else expr

    def visit_node(sub_ast):
        if isinstance(sub_ast, ast.SympyAssignment):
            enclosing_block = sub_ast.parent
            assert type(enclosing_block) is ast.Block
            sub_ast.lhs = visit_sympy_expr(sub_ast.lhs, enclosing_block, sub_ast)
            sub_ast.rhs = visit_sympy_expr(sub_ast.rhs, enclosing_block, sub_ast)
446
        else:
Martin Bauer's avatar
Martin Bauer committed
447
448
            for i, a in enumerate(sub_ast.args):
                visit_node(a)
449

Martin Bauer's avatar
Martin Bauer committed
450
    return visit_node(ast_node)
451
452


Martin Bauer's avatar
Martin Bauer committed
453
def move_constants_before_loop(ast_node):
454
455
    """Moves :class:`pystencils.ast.SympyAssignment` nodes out of loop body if they are iteration independent.

Martin Bauer's avatar
Martin Bauer committed
456
    Call this after creating the loop structure with :func:`make_loop_over_domain`
Martin Bauer's avatar
Martin Bauer committed
457
    """
Martin Bauer's avatar
Martin Bauer committed
458
    def find_block_to_move_to(node):
Martin Bauer's avatar
Martin Bauer committed
459
460
        """
        Traverses parents of node as long as the symbols are independent and returns a (parent) block
461
        the assignment can be safely moved to
Martin Bauer's avatar
Martin Bauer committed
462
463
464
        :param node: SympyAssignment inside a Block
        :return blockToInsertTo, childOfBlockToInsertBefore
        """
465
466
        assert isinstance(node.parent, ast.Block)

Martin Bauer's avatar
Martin Bauer committed
467
468
        last_block = node.parent
        last_block_child = node
469
        element = node.parent
Martin Bauer's avatar
Martin Bauer committed
470
        prev_element = node
471
472
        while element:
            if isinstance(element, ast.Block):
Martin Bauer's avatar
Martin Bauer committed
473
474
                last_block = element
                last_block_child = prev_element
475
476

            if isinstance(element, ast.Conditional):
Martin Bauer's avatar
Martin Bauer committed
477
                critical_symbols = element.condition_expr.atoms(sp.Symbol)
478
            else:
Martin Bauer's avatar
Martin Bauer committed
479
480
                critical_symbols = element.symbols_defined
            if node.undefined_symbols.intersection(critical_symbols):
481
                break
Martin Bauer's avatar
Martin Bauer committed
482
            prev_element = element
483
            element = element.parent
Martin Bauer's avatar
Martin Bauer committed
484
        return last_block, last_block_child
485

Martin Bauer's avatar
Martin Bauer committed
486
487
    def check_if_assignment_already_in_block(assignment, target_block):
        for arg in target_block.args:
488
489
490
491
492
493
            if type(arg) is not ast.SympyAssignment:
                continue
            if arg.lhs == assignment.lhs:
                return arg
        return None

Martin Bauer's avatar
Martin Bauer committed
494
    def get_blocks(node, result_list):
495
        if isinstance(node, ast.Block):
Martin Bauer's avatar
Martin Bauer committed
496
            result_list.append(node)
497
498
        if isinstance(node, ast.Node):
            for a in node.args:
Martin Bauer's avatar
Martin Bauer committed
499
                get_blocks(a, result_list)
500

Martin Bauer's avatar
Martin Bauer committed
501
502
503
    all_blocks = []
    get_blocks(ast_node, all_blocks)
    for block in all_blocks:
Martin Bauer's avatar
Martin Bauer committed
504
        children = block.take_child_nodes()
505
        for child in children:
506
507
508
            target, child_to_insert_before = find_block_to_move_to(child)
            if target == block:     # movement not possible
                target.append(child)
509
            else:
510
511
                if isinstance(child, ast.SympyAssignment):
                    exists_already = check_if_assignment_already_in_block(child, target)
512
                else:
513
                    exists_already = False
Martin Bauer's avatar
Martin Bauer committed
514

515
516
                if not exists_already:
                    target.insert_before(child, child_to_insert_before)
Martin Bauer's avatar
Martin Bauer committed
517
518
                elif exists_already and exists_already.rhs == child.rhs:
                    pass
519
                else:
Martin Bauer's avatar
Martin Bauer committed
520
                    block.append(child)  # don't move in this case - better would be to rename symbol
521
522


Martin Bauer's avatar
Martin Bauer committed
523
def split_inner_loop(ast_node: ast.Node, symbol_groups):
Martin Bauer's avatar
Martin Bauer committed
524
525
    """
    Splits inner loop into multiple loops to minimize the amount of simultaneous load/store streams
526

Martin Bauer's avatar
Martin Bauer committed
527
528
529
530
531
    Args:
        ast_node: AST root
        symbol_groups: sequence of symbol sequences: for each symbol sequence a new inner loop is created which
                       updates these symbols and their dependent symbols. Symbols which are in none of the symbolGroups
                       and which no symbol in a symbol group depends on, are not updated!
Martin Bauer's avatar
Martin Bauer committed
532
    """
Martin Bauer's avatar
Martin Bauer committed
533
534
535
536
537
538
539
540
541
542
543
544
545
    all_loops = ast_node.atoms(ast.LoopOverCoordinate)
    inner_loop = [l for l in all_loops if l.is_innermost_loop]
    assert len(inner_loop) == 1, "Error in AST: multiple innermost loops. Was split transformation already called?"
    inner_loop = inner_loop[0]
    assert type(inner_loop.body) is ast.Block
    outer_loop = [l for l in all_loops if l.is_outermost_loop]
    assert len(outer_loop) == 1, "Error in AST, multiple outermost loops."
    outer_loop = outer_loop[0]

    symbols_with_temporary_array = OrderedDict()
    assignment_map = OrderedDict((a.lhs, a) for a in inner_loop.body.args)

    assignment_groups = []
Martin Bauer's avatar
Martin Bauer committed
546
    for symbol_group in symbol_groups:
547
        # get all dependent symbols
Martin Bauer's avatar
Martin Bauer committed
548
        symbols_to_process = list(symbol_group)
Martin Bauer's avatar
Martin Bauer committed
549
550
551
552
        symbols_resolved = set()
        while symbols_to_process:
            s = symbols_to_process.pop()
            if s in symbols_resolved:
553
554
                continue

Martin Bauer's avatar
Martin Bauer committed
555
            if s in assignment_map:  # if there is no assignment inside the loop body it is independent already
Martin Bauer's avatar
Martin Bauer committed
556
557
558
                for new_symbol in assignment_map[s].rhs.atoms(sp.Symbol):
                    if type(new_symbol) is not Field.Access and new_symbol not in symbols_with_temporary_array:
                        symbols_to_process.append(new_symbol)
Martin Bauer's avatar
Martin Bauer committed
559
            symbols_resolved.add(s)
560

Martin Bauer's avatar
Martin Bauer committed
561
        for symbol in symbol_group:
562
563
            if type(symbol) is not Field.Access:
                assert type(symbol) is TypedSymbol
Martin Bauer's avatar
Martin Bauer committed
564
565
                new_ts = TypedSymbol(symbol.name, PointerType(symbol.dtype))
                symbols_with_temporary_array[symbol] = IndexedBase(new_ts, shape=(1,))[inner_loop.loop_counter_symbol]
566

Martin Bauer's avatar
Martin Bauer committed
567
568
569
570
        assignment_group = []
        for assignment in inner_loop.body.args:
            if assignment.lhs in symbols_resolved:
                new_rhs = assignment.rhs.subs(symbols_with_temporary_array.items())
Martin Bauer's avatar
Martin Bauer committed
571
                if type(assignment.lhs) is not Field.Access and assignment.lhs in symbol_group:
572
                    assert type(assignment.lhs) is TypedSymbol
Martin Bauer's avatar
Martin Bauer committed
573
574
                    new_ts = TypedSymbol(assignment.lhs.name, PointerType(assignment.lhs.dtype))
                    new_lhs = IndexedBase(new_ts, shape=(1,))[inner_loop.loop_counter_symbol]
575
                else:
Martin Bauer's avatar
Martin Bauer committed
576
577
578
                    new_lhs = assignment.lhs
                assignment_group.append(ast.SympyAssignment(new_lhs, new_rhs))
        assignment_groups.append(assignment_group)
579

Martin Bauer's avatar
Martin Bauer committed
580
581
    new_loops = [inner_loop.new_loop_with_different_body(ast.Block(group)) for group in assignment_groups]
    inner_loop.parent.replace(inner_loop, ast.Block(new_loops))
582

Martin Bauer's avatar
Martin Bauer committed
583
584
    for tmp_array in symbols_with_temporary_array:
        tmp_array_pointer = TypedSymbol(tmp_array.name, PointerType(tmp_array.dtype))
Martin Bauer's avatar
Martin Bauer committed
585
586
587
588
        alloc_node = ast.TemporaryMemoryAllocation(tmp_array_pointer, inner_loop.stop, inner_loop.start)
        free_node = ast.TemporaryMemoryFree(alloc_node)
        outer_loop.parent.insert_front(alloc_node)
        outer_loop.parent.append(free_node)
589
590


Martin Bauer's avatar
Martin Bauer committed
591
def cut_loop(loop_node, cutting_points):
Martin Bauer's avatar
Martin Bauer committed
592
593
594
595
596
597
598
599
600
601
    """Cuts loop at given cutting points.

    One loop is transformed into len(cuttingPoints)+1 new loops that range from
    old_begin to cutting_points[1], ..., cutting_points[-1] to old_end

    Modifies the ast in place

    Returns:
        list of new loop nodes
    """
Martin Bauer's avatar
Martin Bauer committed
602
    if loop_node.step != 1:
603
        raise NotImplementedError("Can only split loops that have a step of 1")
Martin Bauer's avatar
Martin Bauer committed
604
605
606
    new_loops = []
    new_start = loop_node.start
    cutting_points = list(cutting_points) + [loop_node.stop]
Martin Bauer's avatar
Martin Bauer committed
607
608
    for new_end in cutting_points:
        if new_end - new_start == 1:
Martin Bauer's avatar
Martin Bauer committed
609
610
611
            new_body = deepcopy(loop_node.body)
            new_body.subs({loop_node.loop_counter_symbol: new_start})
            new_loops.append(new_body)
Martin Bauer's avatar
Martin Bauer committed
612
613
        elif new_end - new_start == 0:
            pass
614
        else:
Martin Bauer's avatar
Martin Bauer committed
615
616
            new_loop = ast.LoopOverCoordinate(deepcopy(loop_node.body), loop_node.coordinate_to_loop_over,
                                              new_start, new_end, loop_node.step)
Martin Bauer's avatar
Martin Bauer committed
617
            new_loops.append(new_loop)
Martin Bauer's avatar
Martin Bauer committed
618
        new_start = new_end
Martin Bauer's avatar
Martin Bauer committed
619
    loop_node.parent.replace(loop_node, new_loops)
Martin Bauer's avatar
Martin Bauer committed
620
    return new_loops
621
622


623
624
def simplify_conditionals(node: ast.Node, loop_counter_simplification: bool=False) -> None:
    """Removes conditionals that are always true/false.
625
626

    Args:
627
628
629
630
631
632
        node: ast node, all descendants of this node are simplified
        loop_counter_simplification: if enabled, tries to detect if a conditional is always true/false
                                     depending on the surrounding loop. For example if the surrounding loop goes from
                                     x=0 to 10 and the condition is x < 0, it is removed.
                                     This analysis needs the integer set library (ISL) islpy, so it is not done by
                                     default.
633
    """
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
    for conditional in node.atoms(ast.Conditional):
        conditional.condition_expr = sp.simplify(conditional.condition_expr)
        if conditional.condition_expr == sp.true:
            conditional.parent.replace(conditional, [conditional.true_block])
        elif conditional.condition_expr == sp.false:
            conditional.parent.replace(conditional, [conditional.false_block] if conditional.false_block else [])
        elif loop_counter_simplification:
            try:
                # noinspection PyUnresolvedReferences
                from pystencils.integer_set_analysis import simplify_loop_counter_dependent_conditional
                simplify_loop_counter_dependent_conditional(conditional)
            except ImportError:
                warnings.warn("Integer simplifications in conditionals skipped, because ISLpy package not installed")


def cleanup_blocks(node: ast.Node) -> None:
650
651
652
653
654
    """Curly Brace Removal: Removes empty blocks, and replaces blocks with a single child by its child """
    if isinstance(node, ast.SympyAssignment):
        return
    elif isinstance(node, ast.Block):
        for a in list(node.args):
Martin Bauer's avatar
Martin Bauer committed
655
            cleanup_blocks(a)
656
657
658
659
660
        if len(node.args) <= 1 and isinstance(node.parent, ast.Block):
            node.parent.replace(node, node.args)
            return
    else:
        for a in node.args:
Martin Bauer's avatar
Martin Bauer committed
661
            cleanup_blocks(a)
662
663


664
665
666
667
668
669
670
671
672
673
674
675
class KernelConstraintsCheck:
    """Checks if the input to create_kernel is valid.

    Test the following conditions:

    - SSA Form for pure symbols:
        -  Every pure symbol may occur only once as left-hand-side of an assignment
        -  Every pure symbol that is read, may not be written to later
    - Independence / Parallelization condition:
        - a field that is written may only be read at exact the same spatial position

    (Pure symbols are symbols that are not Field.Accesses)
Martin Bauer's avatar
Martin Bauer committed
676
    """
677
678
679
680
681
    FieldAndIndex = namedtuple('FieldAndIndex', ['field', 'index'])

    def __init__(self, type_for_symbol, check_independence_condition):
        self._type_for_symbol = type_for_symbol

682
        self.scopes = NestedScopes()
683
684
685
686
687
688
689
690
691
692
        self._field_writes = defaultdict(set)
        self.fields_read = set()
        self.check_independence_condition = check_independence_condition

    def process_assignment(self, assignment):
        # for checks it is crucial to process rhs before lhs to catch e.g. a = a + 1
        new_rhs = self.process_expression(assignment.rhs)
        new_lhs = self._process_lhs(assignment.lhs)
        return ast.SympyAssignment(new_lhs, new_rhs)

693
    def process_expression(self, rhs, type_constants=True):
694
695
        self._update_accesses_rhs(rhs)
        if isinstance(rhs, Field.Access):
Martin Bauer's avatar
Martin Bauer committed
696
            self.fields_read.add(rhs.field)
Martin Bauer's avatar
Martin Bauer committed
697
            self.fields_read.update(rhs.indirect_addressing_fields)
698
699
700
701
            return rhs
        elif isinstance(rhs, TypedSymbol):
            return rhs
        elif isinstance(rhs, sp.Symbol):
702
            return TypedSymbol(rhs.name, self._type_for_symbol[rhs.name])
703
        elif type_constants and isinstance(rhs, sp.Number):
704
705
            return cast_func(rhs, create_type(self._type_for_symbol['_constant']))
        elif isinstance(rhs, sp.Mul):
706
            new_args = [self.process_expression(arg, type_constants) if arg not in (-1, 1) else arg for arg in rhs.args]
707
            return rhs.func(*new_args) if new_args else rhs
708
709
        elif isinstance(rhs, sp.Indexed):
            return rhs
710
711
712
        else:
            if isinstance(rhs, sp.Pow):
                # don't process exponents -> they should remain integers
713
                return sp.Pow(self.process_expression(rhs.args[0], type_constants), rhs.args[1])
714
            else:
715
                new_args = [self.process_expression(arg, type_constants) for arg in rhs.args]
716
                return rhs.func(*new_args) if new_args else rhs
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734

    @property
    def fields_written(self):
        return set(k.field for k, v in self._field_writes.items() if len(v))

    def _process_lhs(self, lhs):
        assert isinstance(lhs, sp.Symbol)
        self._update_accesses_lhs(lhs)
        if not isinstance(lhs, Field.Access) and not isinstance(lhs, TypedSymbol):
            return TypedSymbol(lhs.name, self._type_for_symbol[lhs.name])
        else:
            return lhs

    def _update_accesses_lhs(self, lhs):
        if isinstance(lhs, Field.Access):
            fai = self.FieldAndIndex(lhs.field, lhs.index)
            self._field_writes[fai].add(lhs.offsets)
            if len(self._field_writes[fai]) > 1:
735
                raise ValueError("Field {} is written at two different locations".format(lhs.field.name))
736
        elif isinstance(lhs, sp.Symbol):
737
            if self.scopes.is_defined_locally(lhs):
738
                raise ValueError("Assignments not in SSA form, multiple assignments to {}".format(lhs.name))
739
            if lhs in self.scopes.free_parameters:
740
                raise ValueError("Symbol {} is written, after it has been read".format(lhs.name))
741
            self.scopes.define_symbol(lhs)
742
743
744
745
746
747
748

    def _update_accesses_rhs(self, rhs):
        if isinstance(rhs, Field.Access) and self.check_independence_condition:
            writes = self._field_writes[self.FieldAndIndex(rhs.field, rhs.index)]
            for write_offset in writes:
                assert len(writes) == 1
                if write_offset != rhs.offsets:
749
750
                    raise ValueError("Violation of loop independence condition. Field "
                                     "{} is read at {} and written at {}".format(rhs.field, rhs.offsets, write_offset))
751
752
            self.fields_read.add(rhs.field)
        elif isinstance(rhs, sp.Symbol):
753
            self.scopes.access_symbol(rhs)
754
755
756
757
758


def add_types(eqs, type_for_symbol, check_independence_condition):
    """Traverses AST and replaces every :class:`sympy.Symbol` by a :class:`pystencils.typedsymbol.TypedSymbol`.

Martin Bauer's avatar
Martin Bauer committed
759
760
    Additionally returns sets of all fields which are read/written

761
762
763
764
765
766
767
768
769
    Args:
        eqs: list of equations
        type_for_symbol: dict mapping symbol names to types. Types are strings of C types like 'int' or 'double'
        check_independence_condition: check that loop iterations are independent - this has to be skipped for indexed
                                      kernels

    Returns:
        ``fields_read, fields_written, typed_equations`` set of read fields, set of written fields,
         list of equations where symbols have been replaced by typed symbols
Martin Bauer's avatar
Martin Bauer committed
770
    """
Martin Bauer's avatar
Martin Bauer committed
771
772
    if isinstance(type_for_symbol, str) or not hasattr(type_for_symbol, '__getitem__'):
        type_for_symbol = typing_from_sympy_inspection(eqs, type_for_symbol)
773

774
    check = KernelConstraintsCheck(type_for_symbol, check_independence_condition)
775

Martin Bauer's avatar
Martin Bauer committed
776
777
778
779
    def visit(obj):
        if isinstance(obj, list) or isinstance(obj, tuple):
            return [visit(e) for e in obj]
        if isinstance(obj, sp.Eq) or isinstance(obj, ast.SympyAssignment) or isinstance(obj, Assignment):
780
            return check.process_assignment(obj)
Martin Bauer's avatar
Martin Bauer committed
781
        elif isinstance(obj, ast.Conditional):
782
            check.scopes.push()
Martin Bauer's avatar
Martin Bauer committed
783
            false_block = None if obj.false_block is None else visit(obj.false_block)
784
785
786
787
            result = ast.Conditional(check.process_expression(obj.condition_expr, type_constants=False),
                                     true_block=visit(obj.true_block), false_block=false_block)
            check.scopes.pop()
            return result
Martin Bauer's avatar
Martin Bauer committed
788
        elif isinstance(obj, ast.Block):
789
790
791
792
            check.scopes.push()
            result = ast.Block([visit(e) for e in obj.args])
            check.scopes.pop()
            return result
793
        elif isinstance(obj, ast.Node) and not isinstance(obj, ast.LoopOverCoordinate):
Martin Bauer's avatar
Martin Bauer committed
794
            return obj
795
796
        else:
            raise ValueError("Invalid object in kernel " + str(type(obj)))
797

Martin Bauer's avatar
Martin Bauer committed
798
    typed_equations = visit(eqs)
799

800
    return check.fields_read, check.fields_written, typed_equations
801
802


Martin Bauer's avatar
Martin Bauer committed
803
def insert_casts(node):
804
    """Checks the types and inserts casts and pointer arithmetic where necessary.
Martin Bauer's avatar
Martin Bauer committed
805

806
807
808
809
810
    Args:
        node: the head node of the ast

    Returns:
        modified AST
Martin Bauer's avatar
Martin Bauer committed
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
    """
    def cast(zipped_args_types, target_dtype):
        """
        Adds casts to the arguments if their type differs from the target type
        :param zipped_args_types: a zipped list of args and types
        :param target_dtype: The target data type
        :return: args with possible casts
        """
        casted_args = []
        for argument, data_type in zipped_args_types:
            if data_type.numpy_dtype != target_dtype.numpy_dtype:  # ignoring const
                casted_args.append(cast_func(argument, target_dtype))
            else:
                casted_args.append(argument)
        return casted_args

    def pointer_arithmetic(expr_args):
        """
        Creates a valid pointer arithmetic function
        :param expr_args: Arguments of the add expression
        :return: pointer_arithmetic_func
        """
        pointer = None
        new_args = []
        for arg, data_type in expr_args:
            if data_type.func is PointerType:
                assert pointer is None
                pointer = arg
        for arg, data_type in expr_args:
            if arg != pointer:
                assert data_type.is_int() or data_type.is_uint()
                new_args.append(arg)
        new_args = sp.Add(*new_args) if len(new_args) > 0 else new_args
        return pointer_arithmetic_func(pointer, new_args)

846
    if isinstance(node, sp.AtomicExpr) or isinstance(node, cast_func):
Martin Bauer's avatar
Martin Bauer committed
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
        return node
    args = []
    for arg in node.args:
        args.append(insert_casts(arg))
    # TODO indexed, LoopOverCoordinate
    if node.func in (sp.Add, sp.Mul, sp.Or, sp.And, sp.Pow, sp.Eq, sp.Ne, sp.Lt, sp.Le, sp.Gt, sp.Ge):
        # TODO optimize pow, don't cast integer on double
        types = [get_type_of_expression(arg) for arg in args]
        assert len(types) > 0
        target = collate_types(types)
        zipped = list(zip(args, types))
        if target.func is PointerType:
            assert node.func is sp.Add
            return pointer_arithmetic(zipped)
        else:
            return node.func(*cast(zipped, target))
    elif node.func is ast.SympyAssignment:
        lhs = args[0]
        rhs = args[1]
        target = get_type_of_expression(lhs)
        if target.func is PointerType:
            return node.func(*args)  # TODO fix, not complete
        else:
            return node.func(lhs, *cast([(rhs, get_type_of_expression(rhs))], target))
    elif node.func is ast.ResolvedFieldAccess:
        return node
    elif node.func is ast.Block:
        for old_arg, new_arg in zip(node.args, args):
            node.replace(old_arg, new_arg)
        return node
    elif node.func is ast.LoopOverCoordinate:
        for old_arg, new_arg in zip(node.args, args):
            node.replace(old_arg, new_arg)
        return node
    elif node.func is sp.Piecewise:
        expressions = [expr for (expr, _) in args]
        types = [get_type_of_expression(expr) for expr in expressions]
        target = collate_types(types)
        zipped = list(zip(expressions, types))
        casted_expressions = cast(zipped, target)
        args = [arg.func(*[expr, arg.cond]) for (arg, expr) in zip(args, casted_expressions)]

    return node.func(*args)


892
893
894
895
896
897
898
899
def remove_conditionals_in_staggered_kernel(function_node: ast.KernelFunction) -> None:
    """Removes conditionals of a kernel that iterates over staggered positions by splitting the loops at last element"""

    all_inner_loops = [l for l in function_node.atoms(ast.LoopOverCoordinate) if l.is_innermost_loop]
    assert len(all_inner_loops) == 1, "Transformation works only on kernels with exactly one inner loop"
    inner_loop = all_inner_loops.pop()

    for loop in parents_of_type(inner_loop, ast.LoopOverCoordinate, include_current=True):
900
        cut_loop(loop, [loop.stop - 1])
901

902
    simplify_conditionals(function_node.body, loop_counter_simplification=True)
903
    cleanup_blocks(function_node.body)
Martin Bauer's avatar
Martin Bauer committed
904

905
906
907
908
    move_constants_before_loop(function_node.body)
    cleanup_blocks(function_node.body)


Martin Bauer's avatar
Martin Bauer committed
909
910
911
# --------------------------------------- Helper Functions -------------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
912
def typing_from_sympy_inspection(eqs, default_type="double"):
Martin Bauer's avatar
Martin Bauer committed
913
914
915
    """
    Creates a default symbol name to type mapping.
    If a sympy Boolean is assigned to a symbol it is assumed to be 'bool' otherwise the default type, usually ('double')
916
917
918
919
920
921

    Args:
        eqs: list of equations
        default_type: the type for non-boolean symbols
    Returns:
        dictionary, mapping symbol name to type
Martin Bauer's avatar
Martin Bauer committed
922
    """
Martin Bauer's avatar
Martin Bauer committed
923
    result = defaultdict(lambda: default_type)
924
    for eq in eqs:
925
926
        if isinstance(eq, ast.Node):
            continue
927
928
929
        # problematic case here is when rhs is a symbol: then it is impossible to decide here without
        # further information what type the left hand side is - default fallback is the dict value then
        if isinstance(eq.rhs, Boolean) and not isinstance(eq.rhs, sp.Symbol):
930
931
932
933
            result[eq.lhs.name] = "bool"
    return result


Martin Bauer's avatar
Martin Bauer committed
934
def get_next_parent_of_type(node, parent_type):
935
936
937
938
    """Returns the next parent node of given type or None, if root is reached.

    Traverses the AST nodes parents until a parent of given type was found.
    If no such parent is found, None is returned
Martin Bauer's avatar
Martin Bauer committed
939
    """
940
941
    parent = node.parent
    while parent is not None:
Martin Bauer's avatar
Martin Bauer committed
942
        if isinstance(parent, parent_type):
943
944
945
946
947
            return parent
        parent = parent.parent
    return None


948
def parents_of_type(node, parent_type, include_current=False):
949
    """Generator for all parent nodes of given type"""
950
951
952
953
954
955
956
    parent = node if include_current else node.parent
    while parent is not None:
        if isinstance(parent, parent_type):
            yield parent
        parent = parent.parent


Martin Bauer's avatar
Martin Bauer committed
957
def get_optimal_loop_ordering(fields):
Martin Bauer's avatar
Martin Bauer committed
958
959
960
    """
    Determines the optimal loop order for a given set of fields.
    If the fields have different memory layout or different sizes an exception is thrown.
961
962
963
964
965
966

    Args:
        fields: sequence of fields

    Returns:
        list of coordinate ids, where the first list entry should be the outermost loop
Martin Bauer's avatar
Martin Bauer committed
967
    """
968
    assert len(fields) > 0
Martin Bauer's avatar
Martin Bauer committed
969
    ref_field = next(iter(fields))
970
    for field in fields:
Martin Bauer's avatar
Martin Bauer committed
971
        if field.spatial_dimensions != ref_field.spatial_dimensions:
972
            raise ValueError("All fields have to have the same number of spatial dimensions. Spatial field dimensions: "
Martin Bauer's avatar
Martin Bauer committed
973
                             + str({f.name: f.spatial_shape for f in fields}))
974
975
976

    layouts = set([field.layout for field in fields])
    if len(layouts) > 1:
977
978
        raise ValueError("Due to different layout of the fields no optimal loop ordering exists " +
                         str({f.name: f.layout for f in fields}))
979
    layout = list(layouts)[0]
980
    return list(layout)
981
982


Martin Bauer's avatar
Martin Bauer committed
983
984
985
986
987
988
989
990
991
def replace_inner_stride_with_one(ast_node: ast.KernelFunction) -> None:
    """Replaces the stride of the innermost loop of a variable sized kernel with 1 (assumes optimal loop ordering).

    Variable sized kernels can handle arbitrary field sizes and field shapes. However, the kernel is most efficient
    if the innermost loop accesses the fields with stride 1. The inner loop can also only be vectorized if the inner
    stride is 1. This transformation hard codes this inner stride to one to enable e.g. vectorization.

    Warning: the assumption is not checked at runtime!
    """
992
993