kernelcreation.py 7.4 KB
Newer Older
1
from pystencils.gpucuda.indexing import BlockIndexing
2
from pystencils.transformations import resolve_field_accesses, add_types, parse_base_pointer_info, \
Martin Bauer's avatar
Martin Bauer committed
3
    get_common_shape, resolve_buffer_accesses, unify_shape_symbols, get_base_buffer_index
4
from pystencils.astnodes import Block, KernelFunction, SympyAssignment, LoopOverCoordinate
5
from pystencils.data_types import TypedSymbol, BasicType, StructType
6
from pystencils import Field, FieldType
Martin Bauer's avatar
Martin Bauer committed
7
from pystencils.gpucuda.cudajit import make_python_function
8
9


Martin Bauer's avatar
Martin Bauer committed
10
def create_cuda_kernel(assignments, function_name="kernel", type_info=None, indexing_creator=BlockIndexing,
11
12
                       iteration_slice=None, ghost_layers=None, skip_independence_check=False):
    fields_read, fields_written, assignments = add_types(assignments, type_info, not skip_independence_check)
Martin Bauer's avatar
Martin Bauer committed
13
14
    all_fields = fields_read.union(fields_written)
    read_only_fields = set([f.name for f in fields_read - fields_written])
15

Martin Bauer's avatar
Martin Bauer committed
16
17
    buffers = set([f for f in all_fields if FieldType.is_buffer(f)])
    fields_without_buffers = all_fields - buffers
18

Martin Bauer's avatar
Martin Bauer committed
19
20
21
22
    field_accesses = set()
    num_buffer_accesses = 0
    for eq in assignments:
        field_accesses.update(eq.atoms(Field.Access))
23
        num_buffer_accesses += sum(1 for access in eq.atoms(Field.Access) if FieldType.is_buffer(access.field))
24

Martin Bauer's avatar
Martin Bauer committed
25
    common_shape = get_common_shape(fields_without_buffers)
26

Martin Bauer's avatar
Martin Bauer committed
27
    if iteration_slice is None:
28
        # determine iteration slice from ghost layers
Martin Bauer's avatar
Martin Bauer committed
29
        if ghost_layers is None:
30
            # determine required number of ghost layers from field access
Martin Bauer's avatar
Martin Bauer committed
31
32
33
34
35
36
            required_ghost_layers = max([fa.required_ghost_layers for fa in field_accesses])
            ghost_layers = [(required_ghost_layers, required_ghost_layers)] * len(common_shape)
        iteration_slice = []
        if isinstance(ghost_layers, int):
            for i in range(len(common_shape)):
                iteration_slice.append(slice(ghost_layers, -ghost_layers if ghost_layers > 0 else None))
37
            ghost_layers = [(ghost_layers, ghost_layers)] * len(common_shape)
38
        else:
Martin Bauer's avatar
Martin Bauer committed
39
            for i in range(len(common_shape)):
Martin Bauer's avatar
Martin Bauer committed
40
41
                iteration_slice.append(slice(ghost_layers[i][0],
                                             -ghost_layers[i][1] if ghost_layers[i][1] > 0 else None))
42

Martin Bauer's avatar
Martin Bauer committed
43
    indexing = indexing_creator(field=list(fields_without_buffers)[0], iteration_slice=iteration_slice)
44
45
46
47
48
49
    coord_mapping = indexing.coordinates

    cell_idx_assignments = [SympyAssignment(LoopOverCoordinate.get_loop_counter_symbol(i), value)
                            for i, value in enumerate(coord_mapping)]
    cell_idx_symbols = [LoopOverCoordinate.get_loop_counter_symbol(i) for i, _ in enumerate(coord_mapping)]
    assignments = cell_idx_assignments + assignments
Martin Bauer's avatar
Martin Bauer committed
50

51
    block = Block(assignments)
52

Martin Bauer's avatar
Martin Bauer committed
53
    block = indexing.guard(block, common_shape)
54
55
    unify_shape_symbols(block, common_shape=common_shape, fields=fields_without_buffers)

56
    ast = KernelFunction(block, 'gpu', 'gpucuda', make_python_function, ghost_layers, function_name)
Martin Bauer's avatar
Martin Bauer committed
57
    ast.global_variables.update(indexing.index_variables)
58

Martin Bauer's avatar
Martin Bauer committed
59
60
61
62
    base_pointer_spec = [['spatialInner0']]
    base_pointer_info = {f.name: parse_base_pointer_info(base_pointer_spec, [2, 1, 0],
                                                         f.spatial_dimensions, f.index_dimensions)
                         for f in all_fields}
63

64
    coord_mapping = {f.name: cell_idx_symbols for f in all_fields}
65

Martin Bauer's avatar
Martin Bauer committed
66
    loop_strides = list(fields_without_buffers)[0].shape
67

Martin Bauer's avatar
Martin Bauer committed
68
69
    if any(FieldType.is_buffer(f) for f in all_fields):
        resolve_buffer_accesses(ast, get_base_buffer_index(ast, indexing.coordinates, loop_strides), read_only_fields)
70

Martin Bauer's avatar
Martin Bauer committed
71
    resolve_field_accesses(ast, read_only_fields, field_to_base_pointer_info=base_pointer_info,
Martin Bauer's avatar
Martin Bauer committed
72
                           field_to_fixed_coordinates=coord_mapping)
73

Martin Bauer's avatar
Martin Bauer committed
74
75
    # add the function which determines #blocks and #threads as additional member to KernelFunction node
    # this is used by the jit
76
77
78

    # If loop counter symbols have been explicitly used in the update equations (e.g. for built in periodicity),
    # they are defined here
Martin Bauer's avatar
Martin Bauer committed
79
80
    undefined_loop_counters = {LoopOverCoordinate.is_loop_counter_symbol(s): s for s in ast.body.undefined_symbols
                               if LoopOverCoordinate.is_loop_counter_symbol(s) is not None}
Martin Bauer's avatar
Martin Bauer committed
81
82
    for i, loop_counter in undefined_loop_counters.items():
        ast.body.insert_front(SympyAssignment(loop_counter, indexing.coordinates[i]))
83

84
    ast.indexing = indexing
85
86
87
    return ast


Martin Bauer's avatar
Martin Bauer committed
88
89
def created_indexed_cuda_kernel(assignments, index_fields, function_name="kernel", type_info=None,
                                coordinate_names=('x', 'y', 'z'), indexing_creator=BlockIndexing):
90
    fields_read, fields_written, assignments = add_types(assignments, type_info, check_independence_condition=False)
Martin Bauer's avatar
Martin Bauer committed
91
92
    all_fields = fields_read.union(fields_written)
    read_only_fields = set([f.name for f in fields_read - fields_written])
93

Martin Bauer's avatar
Martin Bauer committed
94
95
96
97
    for index_field in index_fields:
        index_field.field_type = FieldType.INDEXED
        assert FieldType.is_indexed(index_field)
        assert index_field.spatial_dimensions == 1, "Index fields have to be 1D"
Martin Bauer's avatar
Martin Bauer committed
98
99
100
101
102
103
104

    non_index_fields = [f for f in all_fields if f not in index_fields]
    spatial_coordinates = {f.spatial_dimensions for f in non_index_fields}
    assert len(spatial_coordinates) == 1, "Non-index fields do not have the same number of spatial coordinates"
    spatial_coordinates = list(spatial_coordinates)[0]

    def get_coordinate_symbol_assignment(name):
105
106
107
        for ind_f in index_fields:
            assert isinstance(ind_f.dtype, StructType), "Index fields have to have a struct data type"
            data_type = ind_f.dtype
Martin Bauer's avatar
Martin Bauer committed
108
            if data_type.has_element(name):
109
                rhs = ind_f[0](name)
Martin Bauer's avatar
Martin Bauer committed
110
                lhs = TypedSymbol(name, BasicType(data_type.get_element_type(name)))
111
112
113
                return SympyAssignment(lhs, rhs)
        raise ValueError("Index %s not found in any of the passed index fields" % (name,))

Martin Bauer's avatar
Martin Bauer committed
114
115
116
    coordinate_symbol_assignments = [get_coordinate_symbol_assignment(n)
                                     for n in coordinate_names[:spatial_coordinates]]
    coordinate_typed_symbols = [eq.lhs for eq in coordinate_symbol_assignments]
117

Martin Bauer's avatar
Martin Bauer committed
118
119
120
    idx_field = list(index_fields)[0]
    indexing = indexing_creator(field=idx_field,
                                iteration_slice=[slice(None, None, None)] * len(idx_field.spatial_shape))
121

Martin Bauer's avatar
Martin Bauer committed
122
123
    function_body = Block(coordinate_symbol_assignments + assignments)
    function_body = indexing.guard(function_body, get_common_shape(index_fields))
124
    ast = KernelFunction(function_body, 'gpu', 'gpucuda', make_python_function, function_name)
Martin Bauer's avatar
Martin Bauer committed
125
    ast.global_variables.update(indexing.index_variables)
126

Martin Bauer's avatar
Martin Bauer committed
127
    coord_mapping = indexing.coordinates
Martin Bauer's avatar
Martin Bauer committed
128
129
130
131
    base_pointer_spec = [['spatialInner0']]
    base_pointer_info = {f.name: parse_base_pointer_info(base_pointer_spec, [2, 1, 0],
                                                         f.spatial_dimensions, f.index_dimensions)
                         for f in all_fields}
132

Martin Bauer's avatar
Martin Bauer committed
133
134
135
    coord_mapping = {f.name: coord_mapping for f in index_fields}
    coord_mapping.update({f.name: coordinate_typed_symbols for f in non_index_fields})
    resolve_field_accesses(ast, read_only_fields, field_to_fixed_coordinates=coord_mapping,
Martin Bauer's avatar
Martin Bauer committed
136
                           field_to_base_pointer_info=base_pointer_info)
137

138
139
    # add the function which determines #blocks and #threads as additional member to KernelFunction node
    # this is used by the jit
140
141
    ast.indexing = indexing
    return ast