kerncraft_interface.py 7.3 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
2
import warnings
from collections import defaultdict
Martin Bauer's avatar
Martin Bauer committed
3
from tempfile import TemporaryDirectory
Martin Bauer's avatar
Martin Bauer committed
4
from typing import Optional
Martin Bauer's avatar
Martin Bauer committed
5
6

import kerncraft
Martin Bauer's avatar
Martin Bauer committed
7
import sympy as sp
8
from kerncraft.kerncraft import KernelCode
Julian Hammer's avatar
Julian Hammer committed
9
10
from kerncraft.machinemodel import MachineModel

Martin Bauer's avatar
Martin Bauer committed
11
12
from pystencils.astnodes import (
    KernelFunction, LoopOverCoordinate, ResolvedFieldAccess, SympyAssignment)
Martin Bauer's avatar
Martin Bauer committed
13
from pystencils.field import get_layout_from_strides
Martin Bauer's avatar
Martin Bauer committed
14
from pystencils.kerncraft_coupling.generate_benchmark import generate_benchmark
Martin Bauer's avatar
Martin Bauer committed
15
from pystencils.sympyextensions import count_operations_in_ast
16
from pystencils.transformations import filtered_tree_iteration
Martin Bauer's avatar
Martin Bauer committed
17
18
19
from pystencils.utils import DotDict


20
class PyStencilsKerncraftKernel(KernelCode):
Martin Bauer's avatar
Martin Bauer committed
21
22
23
24
    """
    Implementation of kerncraft's kernel interface for pystencils CPU kernels.
    Analyses a list of equations assuming they will be executed on a CPU
    """
Martin Bauer's avatar
Martin Bauer committed
25
26
    LIKWID_BASE = '/usr/local/likwid'

27
28
    def __init__(self, ast: KernelFunction, machine: Optional[MachineModel] = None,
                 assumed_layout='SoA', debug_print=False, filename=None):
Julian Hammer's avatar
Julian Hammer committed
29
30
31
32
33
        """Create a kerncraft kernel using a pystencils AST

        Args:
            ast: pystencils ast
            machine: kerncraft machine model - specify this if kernel needs to be compiled
34
35
36
            assumed_layout: either 'SoA' or 'AoS' - if fields have symbolic sizes the layout of the index
                    coordinates is not known. In this case either a structures of array (SoA) or
                    array of structures (AoS) layout is assumed
Julian Hammer's avatar
Julian Hammer committed
37
        """
Julian Hammer's avatar
Julian Hammer committed
38
        kerncraft.kernel.Kernel.__init__(self, machine)
Martin Bauer's avatar
Martin Bauer committed
39

Julian Hammer's avatar
Julian Hammer committed
40
41
        # Initialize state
        self.asm_block = None
42
        self._filename = filename
Julian Hammer's avatar
Julian Hammer committed
43
44

        self.kernel_ast = ast
Martin Bauer's avatar
Martin Bauer committed
45
        self.temporary_dir = TemporaryDirectory()
46
        self._keep_intermediates = debug_print
Martin Bauer's avatar
Martin Bauer committed
47
48

        # Loops
49
50
        inner_loops = [l for l in filtered_tree_iteration(ast, LoopOverCoordinate, stop_type=SympyAssignment)
                       if l.is_innermost_loop]
Martin Bauer's avatar
Martin Bauer committed
51
        if len(inner_loops) == 0:
Martin Bauer's avatar
Martin Bauer committed
52
53
            raise ValueError("No loop found in pystencils AST")
        else:
54
55
56
            if len(inner_loops) > 1:
                warnings.warn("pystencils AST contains multiple inner loops. "
                              "Only one can be analyzed - choosing first one")
Martin Bauer's avatar
Martin Bauer committed
57
            inner_loop = inner_loops[0]
Martin Bauer's avatar
Martin Bauer committed
58
59

        self._loop_stack = []
Martin Bauer's avatar
Martin Bauer committed
60
61
62
        cur_node = inner_loop
        while cur_node is not None:
            if isinstance(cur_node, LoopOverCoordinate):
Martin Bauer's avatar
Martin Bauer committed
63
                loop_counter_sym = cur_node.loop_counter_symbol
64
                loop_info = (loop_counter_sym.name, cur_node.start, cur_node.stop, 1)
Julian Hammer's avatar
Julian Hammer committed
65
66
                # If the correct step were to be provided, all access within that step length will
                # also need to be passed to kerncraft: cur_node.step)
Martin Bauer's avatar
Martin Bauer committed
67
                self._loop_stack.append(loop_info)
Martin Bauer's avatar
Martin Bauer committed
68
            cur_node = cur_node.parent
Martin Bauer's avatar
Martin Bauer committed
69
70
71
        self._loop_stack = list(reversed(self._loop_stack))

        # Data sources & destinations
72
73
        self.sources = defaultdict(list)
        self.destinations = defaultdict(list)
Martin Bauer's avatar
Martin Bauer committed
74

Julian Hammer's avatar
Julian Hammer committed
75
76
77
78
79
80
81
82
83
        def get_layout_tuple(f):
            if f.has_fixed_shape:
                return get_layout_from_strides(f.strides)
            else:
                layout_list = list(f.layout)
                for _ in range(f.index_dimensions):
                    layout_list.insert(0 if assumed_layout == 'SoA' else -1, max(layout_list) + 1)
                return layout_list

Martin Bauer's avatar
Martin Bauer committed
84
        reads, writes = search_resolved_field_accesses_in_ast(inner_loop)
Martin Bauer's avatar
Martin Bauer committed
85
        for accesses, target_dict in [(reads, self.sources), (writes, self.destinations)]:
Martin Bauer's avatar
Martin Bauer committed
86
            for fa in accesses:
Martin Bauer's avatar
Martin Bauer committed
87
                coord = [sp.Symbol(LoopOverCoordinate.get_loop_counter_name(i), positive=True, integer=True) + off
Martin Bauer's avatar
Martin Bauer committed
88
                         for i, off in enumerate(fa.offsets)]
Martin Bauer's avatar
Martin Bauer committed
89
                coord += list(fa.idx_coordinate_values)
Julian Hammer's avatar
Julian Hammer committed
90
91
                layout = get_layout_tuple(fa.field)
                permuted_coord = [sp.sympify(coord[i]) for i in layout]
Martin Bauer's avatar
Martin Bauer committed
92
                target_dict[fa.field.name].append(permuted_coord)
Martin Bauer's avatar
Martin Bauer committed
93
94

        # Variables (arrays)
Martin Bauer's avatar
Martin Bauer committed
95
96
        fields_accessed = ast.fields_accessed
        for field in fields_accessed:
Julian Hammer's avatar
Julian Hammer committed
97
            layout = get_layout_tuple(field)
Martin Bauer's avatar
Martin Bauer committed
98
99
            permuted_shape = list(field.shape[i] for i in layout)
            self.set_variable(field.name, str(field.dtype), tuple(permuted_shape))
Martin Bauer's avatar
Martin Bauer committed
100

101
102
103
104
105
        # Scalars may be safely ignored
        # for param in ast.get_parameters():
        #     if not param.is_field_parameter:
        #         # self.set_variable(param.symbol.name, str(param.symbol.dtype), None)
        #         self.sources[param.symbol.name] = [None]
Martin Bauer's avatar
Martin Bauer committed
106
107
108
109
110

        # data type
        self.datatype = list(self.variables.values())[0][0]

        # flops
Martin Bauer's avatar
Martin Bauer committed
111
        operation_count = count_operations_in_ast(inner_loop)
Martin Bauer's avatar
Martin Bauer committed
112
        self._flops = {
Martin Bauer's avatar
Martin Bauer committed
113
114
115
            '+': operation_count['adds'],
            '*': operation_count['muls'],
            '/': operation_count['divs'],
Martin Bauer's avatar
Martin Bauer committed
116
        }
Jan Hönig's avatar
Jan Hönig committed
117
118
        for k in [k for k, v in self._flops.items() if v == 0]:
            del self._flops[k]
Martin Bauer's avatar
Martin Bauer committed
119
120
        self.check()

121
122
123
124
125
126
127
128
129
130
131
        if debug_print:
            from pprint import pprint
            print("-----------------------------  Loop Stack --------------------------")
            pprint(self._loop_stack)
            print("-----------------------------  Sources -----------------------------")
            pprint(self.sources)
            print("-----------------------------  Destinations ------------------------")
            pprint(self.destinations)
            print("-----------------------------  FLOPS -------------------------------")
            pprint(self._flops)

132
    def as_code(self, type_='iaca', openmp=False, as_filename=False):
Julian Hammer's avatar
Julian Hammer committed
133
134
135
        """
        Generate and return compilable source code.

136
137
138
139
        Args:
            type_: can be iaca or likwid.
            openmp: if true, openmp code will be generated
            as_filename:
Julian Hammer's avatar
Julian Hammer committed
140
        """
141
142
143
144
145
146
147
148
149
        code = generate_benchmark(self.kernel_ast, likwid=type_ == 'likwid', openmp=openmp)
        if as_filename:
            fp, already_available = self._get_intermediate_file('kernel_{}.c'.format(type_),
                                                                machine_and_compiler_dependent=False)
            if not already_available:
                fp.write(code)
            return fp.name
        else:
            return code
150

Martin Bauer's avatar
Martin Bauer committed
151
152

class KerncraftParameters(DotDict):
Martin Bauer's avatar
Martin Bauer committed
153
154
    def __init__(self, **kwargs):
        super(KerncraftParameters, self).__init__(**kwargs)
Martin Bauer's avatar
Martin Bauer committed
155
156
157
158
159
        self['asm_block'] = 'auto'
        self['asm_increment'] = 0
        self['cores'] = 1
        self['cache_predictor'] = 'SIM'
        self['verbose'] = 0
Jan Hönig's avatar
Jan Hönig committed
160
        self['pointer_increment'] = 'auto'
Jan Hönig's avatar
Jan Hönig committed
161
        self['iterations'] = 10
Julian Hammer's avatar
Julian Hammer committed
162
163
        self['unit'] = 'cy/CL'
        self['ignore_warnings'] = True
Jan Hönig's avatar
Jan Hönig committed
164

Martin Bauer's avatar
Martin Bauer committed
165
166
167
168

# ------------------------------------------- Helper functions ---------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
169
def search_resolved_field_accesses_in_ast(ast):
Martin Bauer's avatar
Martin Bauer committed
170
171
172
173
174
175
176
177
178
    def visit(node, reads, writes):
        if not isinstance(node, SympyAssignment):
            for a in node.args:
                visit(a, reads, writes)
            return

        for expr, accesses in [(node.lhs, writes), (node.rhs, reads)]:
            accesses.update(expr.atoms(ResolvedFieldAccess))

Martin Bauer's avatar
Martin Bauer committed
179
180
181
    read_accesses = set()
    write_accesses = set()
    visit(ast, read_accesses, write_accesses)
Martin Bauer's avatar
Martin Bauer committed
182
    return read_accesses, write_accesses